Add custom auth headers support

.github/workflows/check.yml

@@ -0,0 +1,36 @@
+name: Format/Lint Pipeline
+
+permissions:
+  contents: read
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  style-check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version-file: ".python-version"
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: Install the project
+        run: uv sync --locked --all-extras --dev
+
+      - name: Check format
+        run: uv run ruff format --check --diff
+
+      # TODO: Enable linting
+      # - name: Lint code
+      #   run: uv run ruff check

examples/agentic_mouse_action.py

@@ -16,8 +16,7 @@ async def main() -> None:
                 "viewport": {
                     "width": 1280,
                     "height": 720,
-                }
-
+                },
             },
             fallback_operator_query="click on the search box",
         )
@@ -30,7 +29,7 @@ async def main() -> None:
                 "viewport": {
                     "width": 1280,
                     "height": 720,
-                }
+                },
             },
             fallback_operator_query='type "Narada AI" in the search box',
         )
@@ -43,7 +42,7 @@ async def main() -> None:
                 "viewport": {
                     "width": 1280,
                     "height": 720,
-                }
+                },
             },
             fallback_operator_query="scroll down the page",
         )

packages/narada/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "narada"
-version = "0.1.33a4"
+version = "0.1.33a5"
 description = "Python client SDK for Narada"
 license = "Apache-2.0"
 readme = "README.md"

packages/narada/src/narada/client.py

@@ -56,14 +56,23 @@ class Narada:
     _EXTENSION_UNAUTHENTICATED_INDICATOR_SELECTOR = "#narada-extension-unauthenticated"
     _INITIALIZATION_ERROR_INDICATOR_SELECTOR = "#narada-initialization-error"
 
-    _api_key: str
+    _auth_headers: dict[str, str]
     _console: Console
     _playwright_context_manager: PlaywrightContextManager | None = None
     _playwright: Playwright | None = None
     _cloud_windows: set[CloudBrowserWindow]
 
-    def __init__(self, *, api_key: str | None = None) -> None:
-        self._api_key = api_key or os.environ["NARADA_API_KEY"]
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        auth_headers: dict[str, str] | None = None,
+    ) -> None:
+        if auth_headers is not None:
+            self._auth_headers = auth_headers
+        else:
+            api_key = api_key or os.environ["NARADA_API_KEY"]
+            self._auth_headers = {"x-api-key": api_key}
         self._console = Console()
         self._cloud_windows = set()
 
@@ -93,9 +102,7 @@ async def _fetch_sdk_config(self) -> _SdkConfig | None:
 
         try:
             async with aiohttp.ClientSession() as session:
-                async with session.get(
-                    url, headers={"x-api-key": self._api_key}
-                ) as resp:
+                async with session.get(url, headers=self._auth_headers) as resp:
                     if not resp.ok:
                         logging.warning(
                             "Failed to fetch SDK config: %s %s",
@@ -138,7 +145,7 @@ async def open_and_initialize_browser_window(
         await self._fix_download_behavior(side_panel_page)
 
         return LocalBrowserWindow(
-            api_key=self._api_key,
+            auth_headers=self._auth_headers,
             browser_process_id=launch_browser_result.browser_process_id,
             browser_window_id=browser_window_id,
             config=config,
@@ -171,7 +178,7 @@ async def open_and_initialize_cloud_browser_window(
         async with aiohttp.ClientSession() as session:
             async with session.post(
                 endpoint_url,
-                headers={"x-api-key": self._api_key},
+                headers=self._auth_headers,
                 json=request_body,
                 timeout=aiohttp.ClientTimeout(
                     total=180
@@ -201,7 +208,7 @@ async def open_and_initialize_cloud_browser_window(
                 async with aiohttp.ClientSession() as cleanup_session:
                     async with cleanup_session.post(
                         f"{base_url}/cloud-browser/stop-cloud-browser-session",
-                        headers={"x-api-key": self._api_key},
+                        headers=self._auth_headers,
                         json={"session_id": session_id},
                         timeout=aiohttp.ClientTimeout(total=10),
                     ) as resp:
@@ -249,7 +256,7 @@ async def open_and_initialize_cloud_browser_window(
         cloud_window = CloudBrowserWindow(
             browser_window_id=browser_window_id,
             session_id=session_id,
-            api_key=self._api_key,
+            auth_headers=self._auth_headers,
         )
 
         # Track the window for cleanup in __aexit__
@@ -312,7 +319,7 @@ async def initialize_in_existing_browser_window(
             self._print_success_message(browser_window_id)
 
         return LocalBrowserWindow(
-            api_key=self._api_key,
+            auth_headers=self._auth_headers,
             browser_process_id=None,
             browser_window_id=browser_window_id,
             config=config,

packages/narada/src/narada/window.py

@@ -8,7 +8,6 @@
 from typing import IO, Any, TypeVar, overload
 
 import aiohttp
-from narada.config import BrowserConfig
 from narada_core.actions.models import (
     AgenticMouseAction,
     AgenticMouseActionRequest,
@@ -53,6 +52,8 @@
 )
 from pydantic import BaseModel
 
+from narada.config import BrowserConfig
+
 logger = logging.getLogger(__name__)
 
 _StructuredOutput = TypeVar("_StructuredOutput", bound=BaseModel)
@@ -67,18 +68,18 @@ class _PresignedPost(BaseModel):
 
 
 class BaseBrowserWindow(ABC):
-    _api_key: str
+    _auth_headers: dict[str, str]
     _base_url: str
     _browser_window_id: str
 
     def __init__(
         self,
         *,
-        api_key: str,
+        auth_headers: dict[str, str],
         base_url: str,
         browser_window_id: str,
     ) -> None:
-        self._api_key = api_key
+        self._auth_headers = auth_headers
         self._base_url = base_url
         self._browser_window_id = browser_window_id
 
@@ -99,7 +100,7 @@ async def upload_file(self, *, file: IO) -> File:
             # First generate a presigned POST for uploading the file.
             async with session.post(
                 f"{self._base_url}/remote-dispatch/generate-file-upload-presigned-post",
-                headers={"x-api-key": self._api_key},
+                headers=self._auth_headers,
                 json={"filename": filename},
             ) as resp:
                 resp.raise_for_status()
@@ -192,8 +193,6 @@ async def dispatch_request(
         """
         deadline = time.monotonic() + timeout
 
-        headers = {"x-api-key": self._api_key}
-
         agent_prefix = (
             agent.prompt_prefix() if isinstance(agent, Agent) else f"{agent} "
         )
@@ -238,7 +237,7 @@ async def dispatch_request(
             async with aiohttp.ClientSession() as session:
                 async with session.post(
                     f"{self._base_url}/remote-dispatch",
-                    headers=headers,
+                    headers=self._auth_headers,
                     json=body,
                     timeout=aiohttp.ClientTimeout(total=timeout),
                 ) as resp:
@@ -248,7 +247,7 @@ async def dispatch_request(
                 while (now := time.monotonic()) < deadline:
                     async with session.get(
                         f"{self._base_url}/remote-dispatch/responses/{request_id}",
-                        headers=headers,
+                        headers=self._auth_headers,
                         timeout=aiohttp.ClientTimeout(total=deadline - now),
                     ) as resp:
                         resp.raise_for_status()
@@ -380,6 +379,7 @@ async def agentic_selector(
                 selectors=selectors,
                 fallback_operator_query=fallback_operator_query,
             ),
+            response_model=response_model,
             timeout=timeout,
         )
 
@@ -511,8 +511,6 @@ async def _run_extension_action(
         *,
         timeout: int | None = None,
     ) -> _ResponseModel | None:
-        headers = {"x-api-key": self._api_key}
-
         body = {
             "action": request.model_dump(),
             "browserWindowId": self.browser_window_id,
@@ -523,7 +521,7 @@ async def _run_extension_action(
         async with aiohttp.ClientSession() as session:
             async with session.post(
                 f"{self._base_url}/extension-actions",
-                headers=headers,
+                headers=self._auth_headers,
                 json=body,
                 # Don't specify `timeout` here as the (soft) timeout is handled by the server.
             ) as resp:
@@ -551,15 +549,15 @@ class LocalBrowserWindow(BaseBrowserWindow):
     def __init__(
         self,
         *,
-        api_key: str,
+        auth_headers: dict[str, str],
         browser_process_id: int | None,
         browser_window_id: str,
         config: BrowserConfig,
         context: BrowserContext,
     ) -> None:
         base_url = os.getenv("NARADA_API_BASE_URL", "https://api.narada.ai/fast/v2")
         super().__init__(
-            api_key=api_key,
+            auth_headers=auth_headers,
             base_url=base_url,
             browser_window_id=browser_window_id,
         )
@@ -591,10 +589,19 @@ async def reinitialize(self) -> None:
 
 
 class RemoteBrowserWindow(BaseBrowserWindow):
-    def __init__(self, *, browser_window_id: str, api_key: str | None = None) -> None:
+    def __init__(
+        self,
+        *,
+        browser_window_id: str,
+        api_key: str | None = None,
+        auth_headers: dict[str, str] | None = None,
+    ) -> None:
         base_url = os.getenv("NARADA_API_BASE_URL", "https://api.narada.ai/fast/v2")
+        if auth_headers is None:
+            api_key = api_key or os.environ["NARADA_API_KEY"]
+            auth_headers = {"x-api-key": api_key}
         super().__init__(
-            api_key=api_key or os.environ["NARADA_API_KEY"],
+            auth_headers=auth_headers,
             base_url=base_url,
             browser_window_id=browser_window_id,
         )
@@ -616,10 +623,14 @@ def __init__(
         browser_window_id: str,
         session_id: str,
         api_key: str | None = None,
+        auth_headers: dict[str, str] | None = None,
     ) -> None:
         base_url = os.getenv("NARADA_API_BASE_URL", "https://api.narada.ai/fast/v2")
+        if auth_headers is None:
+            api_key = api_key or os.environ["NARADA_API_KEY"]
+            auth_headers = {"x-api-key": api_key}
         super().__init__(
-            api_key=api_key or os.environ["NARADA_API_KEY"],
+            auth_headers=auth_headers,
             base_url=base_url,
             browser_window_id=browser_window_id,
         )
@@ -631,7 +642,7 @@ async def cleanup(self) -> None:
             async with aiohttp.ClientSession() as session:
                 async with session.post(
                     f"{self._base_url}/cloud-browser/stop-cloud-browser-session",
-                    headers={"x-api-key": self._api_key},
+                    headers=self._auth_headers,
                     json={
                         "session_id": self._session_id,
                     },