Allow direct custom headers in client constructors

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

Author: cursoragent

hyperbrowser/client/async_client.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 
 from ..config import ClientConfig
 from ..transport.async_transport import AsyncTransport
@@ -23,9 +23,10 @@ def __init__(
         config: Optional[ClientConfig] = None,
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
         timeout: Optional[int] = 30,
     ):
-        super().__init__(AsyncTransport, config, api_key, base_url)
+        super().__init__(AsyncTransport, config, api_key, base_url, headers)
         self.transport.client.timeout = timeout
         self.sessions = SessionManager(self)
         self.web = WebManager(self)

hyperbrowser/client/base.py

@@ -1,9 +1,9 @@
-from typing import Optional, Type, Union
+import os
+from typing import Dict, Optional, Type, Union
 
 from hyperbrowser.exceptions import HyperbrowserError
 from ..config import ClientConfig
 from ..transport.base import AsyncTransportStrategy, SyncTransportStrategy
-import os
 
 
 class HyperbrowserBase:
@@ -15,7 +15,15 @@ def __init__(
         config: Optional[ClientConfig] = None,
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
     ):
+        if config is not None and any(
+            value is not None for value in (api_key, base_url, headers)
+        ):
+            raise TypeError(
+                "Pass either `config` or `api_key`/`base_url`/`headers`, not both."
+            )
+
         if config is None:
             config = ClientConfig(
                 api_key=(
@@ -30,6 +38,7 @@ def __init__(
                         "HYPERBROWSER_BASE_URL", "https://api.hyperbrowser.ai"
                     )
                 ),
+                headers=headers,
             )
 
         if not config.api_key:

hyperbrowser/client/sync.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 
 from ..config import ClientConfig
 from ..transport.sync import SyncTransport
@@ -23,9 +23,10 @@ def __init__(
         config: Optional[ClientConfig] = None,
         api_key: Optional[str] = None,
         base_url: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
         timeout: Optional[int] = 30,
     ):
-        super().__init__(SyncTransport, config, api_key, base_url)
+        super().__init__(SyncTransport, config, api_key, base_url, headers)
         self.transport.client.timeout = timeout
         self.sessions = SessionManager(self)
         self.web = WebManager(self)

tests/test_custom_headers.py

@@ -1,5 +1,7 @@
 import asyncio
 
+import pytest
+
 from hyperbrowser import AsyncHyperbrowser, Hyperbrowser
 from hyperbrowser.config import ClientConfig
 from hyperbrowser.transport.async_transport import AsyncTransport
@@ -45,6 +47,14 @@ def test_sync_client_config_headers_are_applied_to_transport():
         client.close()
 
 
+def test_sync_client_constructor_headers_are_applied_to_transport():
+    client = Hyperbrowser(api_key="test-key", headers={"X-Team-Trace": "team-2"})
+    try:
+        assert client.transport.client.headers["X-Team-Trace"] == "team-2"
+    finally:
+        client.close()
+
+
 def test_async_client_config_headers_are_applied_to_transport():
     async def run() -> None:
         client = AsyncHyperbrowser(
@@ -56,3 +66,22 @@ async def run() -> None:
             await client.close()
 
     asyncio.run(run())
+
+
+def test_async_client_constructor_headers_are_applied_to_transport():
+    async def run() -> None:
+        client = AsyncHyperbrowser(api_key="test-key", headers={"X-Team-Trace": "team-2"})
+        try:
+            assert client.transport.client.headers["X-Team-Trace"] == "team-2"
+        finally:
+            await client.close()
+
+    asyncio.run(run())
+
+
+def test_client_constructor_rejects_mixed_config_and_direct_args():
+    with pytest.raises(TypeError, match="Pass either `config`"):
+        Hyperbrowser(
+            config=ClientConfig(api_key="test-key"),
+            headers={"X-Team-Trace": "team-1"},
+        )