Support custom client headers in transport configuration

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/client/base.py

@@ -36,7 +36,7 @@ def __init__(
             raise HyperbrowserError("API key must be provided")
 
         self.config = config
-        self.transport = transport(config.api_key)
+        self.transport = transport(config.api_key, headers=config.headers)
 
     def _build_url(self, path: str) -> str:
         return f"{self.config.base_url}/api{path}"

hyperbrowser/config.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Dict, Optional
 import os
 
 from .exceptions import HyperbrowserError
@@ -10,6 +11,7 @@ class ClientConfig:
 
     api_key: str
     base_url: str = "https://api.hyperbrowser.ai"
+    headers: Optional[Dict[str, str]] = None
 
     @classmethod
     def from_env(cls) -> "ClientConfig":

hyperbrowser/transport/async_transport.py

@@ -10,13 +10,14 @@
 class AsyncTransport(AsyncTransportStrategy):
     """Asynchronous transport implementation using httpx"""
 
-    def __init__(self, api_key: str):
-        self.client = httpx.AsyncClient(
-            headers={
-                "x-api-key": api_key,
-                "User-Agent": f"hyperbrowser-python-sdk/{__version__}",
-            }
-        )
+    def __init__(self, api_key: str, headers: Optional[dict] = None):
+        merged_headers = {
+            "x-api-key": api_key,
+            "User-Agent": f"hyperbrowser-python-sdk/{__version__}",
+        }
+        if headers:
+            merged_headers.update(headers)
+        self.client = httpx.AsyncClient(headers=merged_headers)
         self._closed = False
 
     async def close(self) -> None:

hyperbrowser/transport/base.py

@@ -37,7 +37,7 @@ class SyncTransportStrategy(ABC):
     """Abstract base class for synchronous transport implementations"""
 
     @abstractmethod
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, headers: Optional[dict] = None):
         ...
 
     @abstractmethod
@@ -69,7 +69,7 @@ class AsyncTransportStrategy(ABC):
     """Abstract base class for asynchronous transport implementations"""
 
     @abstractmethod
-    def __init__(self, api_key: str):
+    def __init__(self, api_key: str, headers: Optional[dict] = None):
         ...
 
     @abstractmethod

hyperbrowser/transport/sync.py

@@ -9,13 +9,14 @@
 class SyncTransport(SyncTransportStrategy):
     """Synchronous transport implementation using httpx"""
 
-    def __init__(self, api_key: str):
-        self.client = httpx.Client(
-            headers={
-                "x-api-key": api_key,
-                "User-Agent": f"hyperbrowser-python-sdk/{__version__}",
-            }
-        )
+    def __init__(self, api_key: str, headers: Optional[dict] = None):
+        merged_headers = {
+            "x-api-key": api_key,
+            "User-Agent": f"hyperbrowser-python-sdk/{__version__}",
+        }
+        if headers:
+            merged_headers.update(headers)
+        self.client = httpx.Client(headers=merged_headers)
 
     def _handle_response(self, response: httpx.Response) -> APIResponse:
         try:

tests/test_custom_headers.py

@@ -0,0 +1,58 @@
+import asyncio
+
+from hyperbrowser import AsyncHyperbrowser, Hyperbrowser
+from hyperbrowser.config import ClientConfig
+from hyperbrowser.transport.async_transport import AsyncTransport
+from hyperbrowser.transport.sync import SyncTransport
+
+
+def test_sync_transport_accepts_custom_headers():
+    transport = SyncTransport(
+        api_key="test-key",
+        headers={"X-Correlation-Id": "abc123", "User-Agent": "custom-agent"},
+    )
+    try:
+        assert transport.client.headers["x-api-key"] == "test-key"
+        assert transport.client.headers["X-Correlation-Id"] == "abc123"
+        assert transport.client.headers["User-Agent"] == "custom-agent"
+    finally:
+        transport.close()
+
+
+def test_async_transport_accepts_custom_headers():
+    async def run() -> None:
+        transport = AsyncTransport(
+            api_key="test-key",
+            headers={"X-Correlation-Id": "abc123", "User-Agent": "custom-agent"},
+        )
+        try:
+            assert transport.client.headers["x-api-key"] == "test-key"
+            assert transport.client.headers["X-Correlation-Id"] == "abc123"
+            assert transport.client.headers["User-Agent"] == "custom-agent"
+        finally:
+            await transport.close()
+
+    asyncio.run(run())
+
+
+def test_sync_client_config_headers_are_applied_to_transport():
+    client = Hyperbrowser(
+        config=ClientConfig(api_key="test-key", headers={"X-Team-Trace": "team-1"})
+    )
+    try:
+        assert client.transport.client.headers["X-Team-Trace"] == "team-1"
+    finally:
+        client.close()
+
+
+def test_async_client_config_headers_are_applied_to_transport():
+    async def run() -> None:
+        client = AsyncHyperbrowser(
+            config=ClientConfig(api_key="test-key", headers={"X-Team-Trace": "team-1"})
+        )
+        try:
+            assert client.transport.client.headers["X-Team-Trace"] == "team-1"
+        finally:
+            await client.close()
+
+    asyncio.run(run())