feat: add browser routing cache

examples/browser_routing.py

@@ -0,0 +1,27 @@
+"""Example: direct-to-VM browser routing for process exec and raw HTTP."""
+
+from typing import Any, cast
+
+import httpx
+
+from kernel import Kernel, BrowserRoutingConfig
+
+
+def main() -> None:
+    with Kernel(browser_routing=BrowserRoutingConfig(enabled=True, subresources=("process",))) as client:
+        browsers = cast(Any, client.browsers)
+        browser = browsers.create(headless=True)
+        try:
+            browsers.process.exec(browser.session_id, command="uname", args=["-a"])
+
+            response = cast(httpx.Response, browsers.request(browser.session_id, "GET", "https://example.com"))
+            print("status", response.status_code)
+
+            with browsers.stream(browser.session_id, "GET", "https://example.com") as streamed:
+                print("streamed-bytes", len(streamed.read()))
+        finally:
+            browsers.delete_by_id(browser.session_id)
+
+
+if __name__ == "__main__":
+    main()

src/kernel/__init__.py

@@ -16,6 +16,7 @@
     AsyncKernel,
     AsyncStream,
     RequestOptions,
+    BrowserRoutingConfig,
 )
 from ._models import BaseModel
 from ._version import __title__, __version__
@@ -79,6 +80,7 @@
     "RateLimitError",
     "InternalServerError",
     "Timeout",
+    "BrowserRoutingConfig",
     "RequestOptions",
     "Client",
     "AsyncClient",

src/kernel/_client.py

@@ -29,6 +29,12 @@
     SyncAPIClient,
     AsyncAPIClient,
 )
+from .lib.browser_routing.routing import (
+    BrowserRouteCache,
+    BrowserRoutingConfig,
+    strip_direct_vm_auth,
+    rewrite_direct_vm_options,
+)
 
 if TYPE_CHECKING:
     from .resources import (
@@ -64,6 +70,7 @@
     "Transport",
     "ProxiesTypes",
     "RequestOptions",
+    "BrowserRoutingConfig",
     "Kernel",
     "AsyncKernel",
     "Client",
@@ -79,8 +86,10 @@
 class Kernel(SyncAPIClient):
     # client options
     api_key: str
+    browser_route_cache: BrowserRouteCache
 
     _environment: Literal["production", "development"] | NotGiven
+    _browser_routing: BrowserRoutingConfig | None
 
     def __init__(
         self,
@@ -92,6 +101,7 @@ def __init__(
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
         default_query: Mapping[str, object] | None = None,
+        browser_routing: BrowserRoutingConfig | None = None,
         # Configure a custom httpx client.
         # We provide a `DefaultHttpxClient` class that you can pass to retain the default values we use for `limits`, `timeout` & `follow_redirects`.
         # See the [httpx documentation](https://www.python-httpx.org/api/#client) for more details.
@@ -105,6 +115,7 @@ def __init__(
         # outlining your use-case to help us decide if it should be
         # part of our public interface in the future.
         _strict_response_validation: bool = False,
+        _browser_route_cache: BrowserRouteCache | None = None,
     ) -> None:
         """Construct a new synchronous Kernel client instance.
 
@@ -154,6 +165,8 @@ def __init__(
             custom_query=default_query,
             _strict_response_validation=_strict_response_validation,
         )
+        self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
+        self._browser_routing = browser_routing
 
     @cached_property
     def deployments(self) -> DeploymentsResource:
@@ -266,6 +279,15 @@ def default_headers(self) -> dict[str, str | Omit]:
             **self._custom_headers,
         }
 
+    @override
+    def _prepare_options(self, options: Any) -> Any:
+        options = cast(Any, super()._prepare_options(options))
+        return rewrite_direct_vm_options(options, cache=self.browser_route_cache, config=self._browser_routing)
+
+    @override
+    def _prepare_request(self, request: httpx.Request) -> None:
+        strip_direct_vm_auth(request, cache=self.browser_route_cache)
+
     def copy(
         self,
         *,
@@ -279,6 +301,8 @@ def copy(
         set_default_headers: Mapping[str, str] | None = None,
         default_query: Mapping[str, object] | None = None,
         set_default_query: Mapping[str, object] | None = None,
+        browser_routing: BrowserRoutingConfig | None = None,
+        _browser_route_cache: BrowserRouteCache | None = None,
         _extra_kwargs: Mapping[str, Any] = {},
     ) -> Self:
         """
@@ -312,6 +336,8 @@ def copy(
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
             default_headers=headers,
             default_query=params,
+            browser_routing=browser_routing if browser_routing is not None else self._browser_routing,
+            _browser_route_cache=_browser_route_cache or self.browser_route_cache,
             **_extra_kwargs,
         )
 
@@ -356,8 +382,10 @@ def _make_status_error(
 class AsyncKernel(AsyncAPIClient):
     # client options
     api_key: str
+    browser_route_cache: BrowserRouteCache
 
     _environment: Literal["production", "development"] | NotGiven
+    _browser_routing: BrowserRoutingConfig | None
 
     def __init__(
         self,
@@ -369,6 +397,7 @@ def __init__(
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
         default_query: Mapping[str, object] | None = None,
+        browser_routing: BrowserRoutingConfig | None = None,
         # Configure a custom httpx client.
         # We provide a `DefaultAsyncHttpxClient` class that you can pass to retain the default values we use for `limits`, `timeout` & `follow_redirects`.
         # See the [httpx documentation](https://www.python-httpx.org/api/#asyncclient) for more details.
@@ -382,6 +411,7 @@ def __init__(
         # outlining your use-case to help us decide if it should be
         # part of our public interface in the future.
         _strict_response_validation: bool = False,
+        _browser_route_cache: BrowserRouteCache | None = None,
     ) -> None:
         """Construct a new async AsyncKernel client instance.
 
@@ -431,6 +461,8 @@ def __init__(
             custom_query=default_query,
             _strict_response_validation=_strict_response_validation,
         )
+        self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
+        self._browser_routing = browser_routing
 
     @cached_property
     def deployments(self) -> AsyncDeploymentsResource:
@@ -543,6 +575,15 @@ def default_headers(self) -> dict[str, str | Omit]:
             **self._custom_headers,
         }
 
+    @override
+    async def _prepare_options(self, options: Any) -> Any:
+        options = cast(Any, await super()._prepare_options(options))
+        return rewrite_direct_vm_options(options, cache=self.browser_route_cache, config=self._browser_routing)
+
+    @override
+    async def _prepare_request(self, request: httpx.Request) -> None:
+        strip_direct_vm_auth(request, cache=self.browser_route_cache)
+
     def copy(
         self,
         *,
@@ -556,6 +597,8 @@ def copy(
         set_default_headers: Mapping[str, str] | None = None,
         default_query: Mapping[str, object] | None = None,
         set_default_query: Mapping[str, object] | None = None,
+        browser_routing: BrowserRoutingConfig | None = None,
+        _browser_route_cache: BrowserRouteCache | None = None,
         _extra_kwargs: Mapping[str, Any] = {},
     ) -> Self:
         """
@@ -589,6 +632,8 @@ def copy(
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
             default_headers=headers,
             default_query=params,
+            browser_routing=browser_routing if browser_routing is not None else self._browser_routing,
+            _browser_route_cache=_browser_route_cache or self.browser_route_cache,
             **_extra_kwargs,
         )
 

src/kernel/lib/browser_routing/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+__all__: list[str] = []

src/kernel/lib/browser_routing/raw_http.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from typing import IO, Any, Union, Mapping, cast
+from contextlib import contextmanager, asynccontextmanager
+from collections.abc import Iterable, Iterator, AsyncIterator
+
+import httpx
+
+from .util import sanitize_curl_raw_params
+from .routing import BrowserRoute
+from ..._types import Body, Timeout, NotGiven, not_given
+from ..._models import FinalRequestOptions
+
+BrowserRawContent = Union[bytes, bytearray, memoryview, str, IO[bytes], Iterable[bytes]]
+
+
+def request_via_browser_route(
+    parent: Any,
+    route: BrowserRoute,
+    method: str,
+    url: str,
+    *,
+    content: BrowserRawContent | None = None,
+    json: Body | None = None,
+    headers: Mapping[str, str] | None = None,
+    params: Mapping[str, object] | None = None,
+    timeout: float | Timeout | None | NotGiven = not_given,
+) -> httpx.Response:
+    if json is not None and content is not None:
+        raise TypeError("Passing both `json` and `content` is not supported")
+    query: dict[str, object] = {**sanitize_curl_raw_params(params), "url": url, "jwt": route.jwt}
+    options = FinalRequestOptions.construct(
+        method=method.upper(),
+        url=route.base_url.rstrip("/") + "/curl/raw",
+        params=query,
+        headers=headers or {},
+        content=_normalize_binary_content(content),
+        json_data=json,
+        timeout=_normalize_timeout(timeout),
+    )
+    return cast(httpx.Response, parent.request(httpx.Response, options))
+
+
+@contextmanager
+def stream_via_browser_route(
+    parent: Any,
+    route: BrowserRoute,
+    method: str,
+    url: str,
+    *,
+    content: BrowserRawContent | None = None,
+    headers: Mapping[str, str] | None = None,
+    params: Mapping[str, object] | None = None,
+    timeout: float | Timeout | None | NotGiven = not_given,
+) -> Iterator[httpx.Response]:
+    query: dict[str, Any] = sanitize_curl_raw_params(params)
+    query["jwt"] = route.jwt
+    query["url"] = url
+    request_headers = {k: v for k, v in parent.default_headers.items() if isinstance(v, str)}
+    if content is None:
+        request_headers.pop("Content-Type", None)
+    if headers:
+        request_headers.update(headers)
+    request_headers.pop("Authorization", None)
+    effective_timeout = parent.timeout if isinstance(timeout, NotGiven) else timeout
+    with parent._client.stream(
+        method.upper(),
+        route.base_url.rstrip("/") + "/curl/raw",
+        params=query,
+        headers=request_headers,
+        content=_normalize_binary_content(content),
+        timeout=_normalize_timeout(effective_timeout),
+    ) as response:
+        yield response
+
+
+async def async_request_via_browser_route(
+    parent: Any,
+    route: BrowserRoute,
+    method: str,
+    url: str,
+    *,
+    content: BrowserRawContent | None = None,
+    json: Body | None = None,
+    headers: Mapping[str, str] | None = None,
+    params: Mapping[str, object] | None = None,
+    timeout: float | Timeout | None | NotGiven = not_given,
+) -> httpx.Response:
+    if json is not None and content is not None:
+        raise TypeError("Passing both `json` and `content` is not supported")
+    query: dict[str, object] = {**sanitize_curl_raw_params(params), "url": url, "jwt": route.jwt}
+    options = FinalRequestOptions.construct(
+        method=method.upper(),
+        url=route.base_url.rstrip("/") + "/curl/raw",
+        params=query,
+        headers=headers or {},
+        content=_normalize_binary_content(content),
+        json_data=json,
+        timeout=_normalize_timeout(timeout),
+    )
+    return cast(httpx.Response, await parent.request(httpx.Response, options))
+
+
+@asynccontextmanager
+async def async_stream_via_browser_route(
+    parent: Any,
+    route: BrowserRoute,
+    method: str,
+    url: str,
+    *,
+    content: BrowserRawContent | None = None,
+    headers: Mapping[str, str] | None = None,
+    params: Mapping[str, object] | None = None,
+    timeout: float | Timeout | None | NotGiven = not_given,
+) -> AsyncIterator[httpx.Response]:
+    query: dict[str, Any] = sanitize_curl_raw_params(params)
+    query["jwt"] = route.jwt
+    query["url"] = url
+    request_headers = {k: v for k, v in parent.default_headers.items() if isinstance(v, str)}
+    if content is None:
+        request_headers.pop("Content-Type", None)
+    if headers:
+        request_headers.update(headers)
+    request_headers.pop("Authorization", None)
+    effective_timeout = parent.timeout if isinstance(timeout, NotGiven) else timeout
+    async with parent._client.stream(
+        method.upper(),
+        route.base_url.rstrip("/") + "/curl/raw",
+        params=query,
+        headers=request_headers,
+        content=_normalize_binary_content(content),
+        timeout=_normalize_timeout(effective_timeout),
+    ) as response:
+        yield response
+
+
+def _normalize_timeout(timeout: float | Timeout | None | NotGiven) -> float | Timeout | None:
+    return None if isinstance(timeout, NotGiven) else timeout
+
+
+def _normalize_binary_content(content: BrowserRawContent | None) -> bytes | IO[bytes] | Iterable[bytes] | None:
+    if content is None:
+        return None
+    if isinstance(content, str):
+        return content.encode()
+    if isinstance(content, bytearray):
+        return bytes(content)
+    if isinstance(content, memoryview):
+        return content.tobytes()
+    return content