Merge pull request #93 from kernel/raf/browser-scoped-client

feat: add browser routing cache

Author: rgarcia

examples/browser_routing.py

@@ -0,0 +1,25 @@
+"""Example: direct-to-VM browser routing for process exec and raw HTTP."""
+
+from typing import Any, cast
+
+import httpx
+
+from kernel import Kernel
+
+
+def main() -> None:
+    with Kernel() as client:
+        browsers = cast(Any, client.browsers)
+        browser = browsers.create(headless=True)
+        try:
+            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/_client.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Any, Dict, Mapping, cast
+from typing import TYPE_CHECKING, Any, Dict, Type, Mapping, cast
 from typing_extensions import Self, Literal, override
 
 import httpx
@@ -14,13 +14,15 @@
     Omit,
     Timeout,
     NotGiven,
+    ResponseT,
     Transport,
     ProxiesTypes,
     RequestOptions,
     not_given,
 )
 from ._utils import is_given, get_async_library
 from ._compat import cached_property
+from ._models import FinalRequestOptions
 from ._version import __version__
 from ._streaming import Stream as Stream, AsyncStream as AsyncStream
 from ._exceptions import KernelError, APIStatusError
@@ -29,6 +31,15 @@
     SyncAPIClient,
     AsyncAPIClient,
 )
+from .lib.browser_routing.routing import (
+    BrowserRouteCache,
+    BrowserRoutingConfig,
+    strip_direct_vm_auth,
+    rewrite_direct_vm_options,
+    browser_routing_config_from_env,
+    maybe_evict_browser_route_from_response,
+    maybe_populate_browser_route_cache_from_response,
+)
 
 if TYPE_CHECKING:
     from .resources import (
@@ -79,8 +90,10 @@
 class Kernel(SyncAPIClient):
     # client options
     api_key: str
+    browser_route_cache: BrowserRouteCache
 
     _environment: Literal["production", "development"] | NotGiven
+    _browser_routing: BrowserRoutingConfig
 
     def __init__(
         self,
@@ -105,6 +118,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 +168,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_config_from_env()
 
     @cached_property
     def deployments(self) -> DeploymentsResource:
@@ -266,6 +282,37 @@ 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)
+
+    @override
+    def _process_response(
+        self,
+        *,
+        cast_to: Type[ResponseT],
+        options: FinalRequestOptions,
+        response: httpx.Response,
+        stream: bool,
+        stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
+        retries_taken: int = 0,
+    ) -> ResponseT:
+        maybe_populate_browser_route_cache_from_response(response, cache=self.browser_route_cache)
+        maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache)
+        return super()._process_response(
+            cast_to=cast_to,
+            options=options,
+            response=response,
+            stream=stream,
+            stream_cls=stream_cls,
+            retries_taken=retries_taken,
+        )
+
     def copy(
         self,
         *,
@@ -279,6 +326,7 @@ 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_route_cache: BrowserRouteCache | None = None,
         _extra_kwargs: Mapping[str, Any] = {},
     ) -> Self:
         """
@@ -312,6 +360,7 @@ def copy(
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
             default_headers=headers,
             default_query=params,
+            _browser_route_cache=_browser_route_cache or self.browser_route_cache,
             **_extra_kwargs,
         )
 
@@ -356,8 +405,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
 
     def __init__(
         self,
@@ -382,6 +433,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 +483,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_config_from_env()
 
     @cached_property
     def deployments(self) -> AsyncDeploymentsResource:
@@ -543,6 +597,37 @@ 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)
+
+    @override
+    async def _process_response(
+        self,
+        *,
+        cast_to: Type[ResponseT],
+        options: FinalRequestOptions,
+        response: httpx.Response,
+        stream: bool,
+        stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
+        retries_taken: int = 0,
+    ) -> ResponseT:
+        maybe_populate_browser_route_cache_from_response(response, cache=self.browser_route_cache)
+        maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache)
+        return await super()._process_response(
+            cast_to=cast_to,
+            options=options,
+            response=response,
+            stream=stream,
+            stream_cls=stream_cls,
+            retries_taken=retries_taken,
+        )
+
     def copy(
         self,
         *,
@@ -556,6 +641,7 @@ 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_route_cache: BrowserRouteCache | None = None,
         _extra_kwargs: Mapping[str, Any] = {},
     ) -> Self:
         """
@@ -589,6 +675,7 @@ def copy(
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
             default_headers=headers,
             default_query=params,
+            _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