diff --git a/src/kernel/_client.py b/src/kernel/_client.py index 2c06cb04..dc506418 100644 --- a/src/kernel/_client.py +++ b/src/kernel/_client.py @@ -41,8 +41,11 @@ strip_direct_vm_auth, rewrite_direct_vm_options, browser_routing_config_from_env, + is_stale_direct_vm_auth_response, should_retry_stale_direct_vm_auth, + install_stale_direct_vm_auth_eviction, maybe_evict_browser_route_from_response, + install_async_stale_direct_vm_auth_eviction, maybe_populate_browser_route_cache_from_response, ) @@ -203,6 +206,7 @@ def __init__( ) self.browser_route_cache = _browser_route_cache or BrowserRouteCache() self._browser_routing = browser_routing_config_from_env() + install_stale_direct_vm_auth_eviction(self._client, cache=self.browser_route_cache) @cached_property def deployments(self) -> DeploymentsResource: @@ -365,9 +369,11 @@ def _prepare_request(self, request: httpx.Request) -> None: @override def _should_retry(self, response: httpx.Response) -> bool: - if should_retry_stale_direct_vm_auth(response): - maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) - return True + if is_stale_direct_vm_auth_response(response): + # The route was already evicted by the response hook; retry only when + # the body can be rebuilt, otherwise the caller sees the original auth + # failure and a later call goes to the control plane. + return should_retry_stale_direct_vm_auth(response) return super()._should_retry(response) @override @@ -586,6 +592,7 @@ def __init__( ) self.browser_route_cache = _browser_route_cache or BrowserRouteCache() self._browser_routing = browser_routing_config_from_env() + install_async_stale_direct_vm_auth_eviction(self._client, cache=self.browser_route_cache) @cached_property def deployments(self) -> AsyncDeploymentsResource: @@ -748,9 +755,11 @@ async def _prepare_request(self, request: httpx.Request) -> None: @override def _should_retry(self, response: httpx.Response) -> bool: - if should_retry_stale_direct_vm_auth(response): - maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) - return True + if is_stale_direct_vm_auth_response(response): + # The route was already evicted by the response hook; retry only when + # the body can be rebuilt, otherwise the caller sees the original auth + # failure and a later call goes to the control plane. + return should_retry_stale_direct_vm_auth(response) return super()._should_retry(response) @override diff --git a/src/kernel/lib/browser_routing/routing.py b/src/kernel/lib/browser_routing/routing.py index bad5e4ea..df5f6567 100644 --- a/src/kernel/lib/browser_routing/routing.py +++ b/src/kernel/lib/browser_routing/routing.py @@ -32,6 +32,9 @@ class BrowserRoutingConfig: subresources: tuple[str, ...] = field(default_factory=tuple) +_EVICTION_HOOK_CACHE_ATTR = "_kernel_browser_route_cache" + + _BROWSER_ROUTE_CACHEABLE_PATH = re.compile(r"^/(?:v\d+/)?browsers(?:/[^/]+)?/?$") _BROWSER_DELETE_BY_ID_PATH = re.compile(r"^/(?:v\d+/)?browsers/([^/]+)/?$") _BROWSER_POOL_ACQUIRE_PATH = re.compile(r"^/(?:v\d+/)?browser_pools/[^/]+/acquire/?$") @@ -44,7 +47,17 @@ def browser_routing_config_from_env() -> BrowserRoutingConfig: # Path prefixes eligible for direct-to-VM routing. "telemetry/stream" is # the live SSE endpoint (VM); "telemetry/events" is a historical read # served by the control plane (S2) and must NOT be here. - return BrowserRoutingConfig(subresources=("curl", "telemetry/stream", "computer", "playwright", "process")) + return BrowserRoutingConfig( + subresources=( + "curl", + "telemetry/stream", + "computer", + "playwright", + "process", + "fs", + "logs/stream", + ) + ) if raw.strip() == "": return BrowserRoutingConfig() @@ -188,8 +201,124 @@ def is_stale_direct_vm_auth_response(response: httpx.Response) -> bool: return bool(response.request.url.params.get("jwt")) +def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: BrowserRouteCache) -> None: + """Evict stale direct-to-VM routes as soon as the response status is known. + + httpx reads the body of a non-streamed response inside `send()`, so a caller + that only inspects the returned response never learns the status of a 401/403 + whose body read fails — the read error surfaces from `send()` instead and the + dead route would stay cached, wedging every later call for that session. A + response event hook runs after the status is known and before any body is + read, which keeps eviction independent of the body. For a caller-supplied + `http_client`, the hook is installed into that client's `event_hooks` and + prepended so an existing hook cannot pre-empt eviction by reading a failing + body or raising. + """ + hooks = client.event_hooks.setdefault("response", []) + if _has_eviction_hook(hooks, cache): + return + + def evict(response: httpx.Response) -> None: + if is_stale_direct_vm_auth_response(response): + maybe_evict_browser_route_from_response(response, cache=cache) + + setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache) + hooks.insert(0, evict) + + +def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient, *, cache: BrowserRouteCache) -> None: + """Async counterpart of `install_stale_direct_vm_auth_eviction`.""" + hooks = client.event_hooks.setdefault("response", []) + if _has_eviction_hook(hooks, cache): + return + + async def evict(response: httpx.Response) -> None: + if is_stale_direct_vm_auth_response(response): + maybe_evict_browser_route_from_response(response, cache=cache) + + setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache) + hooks.insert(0, evict) + + +def _has_eviction_hook(hooks: list[Any], cache: BrowserRouteCache) -> bool: + # A copied client shares both the httpx client and the route cache, so the + # hook is registered once per cache instead of once per client. + return any(getattr(hook, _EVICTION_HOOK_CACHE_ATTR, None) is cache for hook in hooks) + + def should_retry_stale_direct_vm_auth(response: httpx.Response) -> bool: - return is_stale_direct_vm_auth_response(response) + """Whether a stale direct-to-VM auth failure can be retried on the control plane. + + A retry rebuilds the request from the original options, so it is only safe when + the body can be serialized again byte for byte. Streamed bodies (e.g. a file + object passed to fs.write_file) are consumed by the direct request, so retrying + would send a truncated or empty body to the control plane. + """ + if not is_stale_direct_vm_auth_response(response): + return False + return direct_vm_request_body_is_replayable(response.request) + + +def direct_vm_request_body_is_replayable(request: httpx.Request) -> bool: + try: + _ = request.content + except httpx.RequestNotRead: + pass + else: + # httpx already buffered the body, so rebuilding it yields the same bytes. + return True + + # httpx encodes multipart bodies as a stream of fields it re-renders per attempt. + fields = getattr(request.stream, "fields", None) + if fields is None: + # A streamed body (file object, iterator or async iterator) cannot be replayed. + return False + return all(_multipart_field_is_replayable(field) for field in cast("list[Any]", fields)) + + +def _multipart_field_is_replayable(field: Any) -> bool: + file = getattr(field, "file", None) + if file is None: + # A data field renders from an in-memory value. + return True + if isinstance(file, (bytes, str)): + return True + if getattr(file, "closed", False): + return False + return _rewind_succeeds(file) + + +def _rewind_succeeds(file: Any) -> bool: + """Whether the file field can actually be rewound for another render. + + `seekable()` is not proof: a wrapper can report True and still raise from + `seek()`, which would render the field as an empty part on the retry. The + only reliable check is to perform the rewind httpx would perform. + """ + seek = getattr(file, "seek", None) + if not callable(seek): + return False + + position: object = None + tell = getattr(file, "tell", None) + if callable(tell): + try: + position = tell() + except Exception: + position = None + + try: + seek(0) + except Exception: + return False + + if isinstance(position, int) and position > 0: + try: + seek(position) + except Exception: + # The field is left rewound, which is where httpx renders it from anyway. + pass + return True def _session_id_from_browser_delete_path(path: str) -> str | None: diff --git a/src/kernel/lib/multipart.py b/src/kernel/lib/multipart.py new file mode 100644 index 00000000..611c6d8c --- /dev/null +++ b/src/kernel/lib/multipart.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import Mapping, Sequence, cast + +from .._utils import is_given + +__all__ = ["indexed_multipart_body"] + + +def indexed_multipart_body(body: object) -> dict[str, object]: + """Flatten a multipart body so that array entries carry their index. + + Endpoints that take an array of objects with a file field need each entry's + fields grouped together: `files[0][dest_path]` pairs with the `files[0][file]` + part, while repeated `files[][dest_path]` names cannot be matched back to + their file. The returned mapping is already flat, so the client's generic + multipart serialization passes the names through untouched and every other + endpoint keeps its existing encoding. + """ + flattened: dict[str, object] = {} + if isinstance(body, Mapping): + for key, value in cast(Mapping[object, object], body).items(): + _flatten(str(key), value, flattened) + return flattened + + +def _flatten(key: str, value: object, out: dict[str, object]) -> None: + if not is_given(value): + return + + if isinstance(value, Mapping): + for child_key, child in cast(Mapping[object, object], value).items(): + _flatten(f"{key}[{child_key}]", child, out) + return + + if isinstance(value, (list, tuple)): + for index, child in enumerate(cast(Sequence[object], value)): + _flatten(f"{key}[{index}]", child, out) + return + + out[key] = value diff --git a/src/kernel/resources/browsers/fs/fs.py b/src/kernel/resources/browsers/fs/fs.py index a89d86f2..ab6f572d 100644 --- a/src/kernel/resources/browsers/fs/fs.py +++ b/src/kernel/resources/browsers/fs/fs.py @@ -48,6 +48,7 @@ async_to_custom_streamed_response_wrapper, ) from ...._base_client import make_request_options +from ....lib.multipart import indexed_multipart_body from ....types.browsers import ( f_move_params, f_upload_params, @@ -510,14 +511,19 @@ def upload( raise ValueError(f"Expected a non-empty value for `id_or_name` but received {id_or_name!r}") extra_headers = {"Accept": "*/*", **(extra_headers or {})} body = deepcopy_with_paths({"files": files}, [["files", "", "file"]]) - extracted_files = extract_files(cast(Mapping[str, object], body), paths=[["files", "", "file"]]) + # The remote filesystem pairs each file part with the sibling fields of the + # same array entry, so both halves of the form use indexed names + # (`files[0][file]`, `files[0][dest_path]`). + extracted_files = extract_files( + cast(Mapping[str, object], body), paths=[["files", "", "file"]], array_format="indices" + ) # It should be noted that the actual Content-Type header that will be # sent to the server will contain a `boundary` parameter, e.g. # multipart/form-data; boundary=---abc-- extra_headers["Content-Type"] = "multipart/form-data" return self._post( path_template("/browsers/{id_or_name}/fs/upload", id_or_name=id_or_name), - body=maybe_transform(body, f_upload_params.FUploadParams), + body=indexed_multipart_body(maybe_transform(body, f_upload_params.FUploadParams)), files=extracted_files, options=make_request_options( extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout @@ -1073,14 +1079,19 @@ async def upload( raise ValueError(f"Expected a non-empty value for `id_or_name` but received {id_or_name!r}") extra_headers = {"Accept": "*/*", **(extra_headers or {})} body = deepcopy_with_paths({"files": files}, [["files", "", "file"]]) - extracted_files = extract_files(cast(Mapping[str, object], body), paths=[["files", "", "file"]]) + # The remote filesystem pairs each file part with the sibling fields of the + # same array entry, so both halves of the form use indexed names + # (`files[0][file]`, `files[0][dest_path]`). + extracted_files = extract_files( + cast(Mapping[str, object], body), paths=[["files", "", "file"]], array_format="indices" + ) # It should be noted that the actual Content-Type header that will be # sent to the server will contain a `boundary` parameter, e.g. # multipart/form-data; boundary=---abc-- extra_headers["Content-Type"] = "multipart/form-data" return await self._post( path_template("/browsers/{id_or_name}/fs/upload", id_or_name=id_or_name), - body=await async_maybe_transform(body, f_upload_params.FUploadParams), + body=indexed_multipart_body(await async_maybe_transform(body, f_upload_params.FUploadParams)), files=extracted_files, options=make_request_options( extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout diff --git a/tests/test_browser_routing.py b/tests/test_browser_routing.py index eb47e040..2ab6ad2a 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -1,15 +1,24 @@ from __future__ import annotations +import io import os import asyncio -from typing import Any, AsyncIterator, cast +from typing import Any, Iterator, AsyncIterator, cast +from pathlib import Path from typing_extensions import override import httpx import respx import pytest -from kernel import Kernel, AsyncKernel, InternalServerError +from kernel import ( + Kernel, + AsyncKernel, + APIConnectionError, + AuthenticationError, + InternalServerError, + PermissionDeniedError, +) from kernel.lib.browser_routing.util import jwt_from_cdp_ws_url from kernel.lib.browser_routing.routing import ( BrowserRoute, @@ -37,6 +46,69 @@ def _fake_browser() -> dict[str, object]: } +def _skip_retry_sleep(_self: object, **_kwargs: object) -> None: + return None + + +class _UnseekableFile(io.RawIOBase): + """A file-like upload body that cannot be rewound, e.g. a pipe. + + `claims_seekable` reproduces a wrapper whose `seekable()` says True while + `seek()` still raises, which would otherwise render as an empty part. + """ + + name = "one.txt" + + def __init__(self, content: bytes, *, claims_seekable: bool = False) -> None: + self._content = content + self._claims_seekable = claims_seekable + + @override + def readable(self) -> bool: + return True + + @override + def read(self, size: int = -1) -> bytes: + chunk = self._content if size < 0 else self._content[:size] + self._content = b"" if size < 0 else self._content[size:] + return chunk + + @override + def tell(self) -> int: + return 0 + + @override + def seekable(self) -> bool: + return self._claims_seekable + + @override + def seek(self, _offset: int, _whence: int = 0) -> int: + raise io.UnsupportedOperation("not seekable") + + +class _NoSeekableAttrFile(io.RawIOBase): + """A file-like upload body whose `seek()` raises and reports no seekability.""" + + name = "one.txt" + + def __init__(self, content: bytes) -> None: + self._content = content + + @override + def readable(self) -> bool: + return True + + @override + def read(self, size: int = -1) -> bytes: + chunk = self._content if size < 0 else self._content[:size] + self._content = b"" if size < 0 else self._content[size:] + return chunk + + @override + def seek(self, _offset: int, _whence: int = 0) -> int: + raise OSError("not seekable") + + def _cache_browser(client: Kernel) -> None: route = browser_route_from_browser(_fake_browser()) assert route is not None @@ -398,6 +470,8 @@ def test_browser_routing_config_from_env_defaults(monkeypatch: pytest.MonkeyPatc "computer", "playwright", "process", + "fs", + "logs/stream", ) @@ -407,7 +481,7 @@ def test_direct_vm_routing_allowlist_segment_boundary() -> None: # stream-prefixed-but-different path is not matched. from kernel.lib.browser_routing.routing import _matches_direct_vm_prefix - prefixes = ("curl", "telemetry/stream", "computer", "playwright", "process") + prefixes = ("curl", "telemetry/stream", "computer", "playwright", "process", "fs", "logs/stream") assert _matches_direct_vm_prefix("telemetry/stream", prefixes) is True assert _matches_direct_vm_prefix("telemetry/stream/x", prefixes) is True assert _matches_direct_vm_prefix("telemetry/events", prefixes) is False @@ -418,7 +492,16 @@ def test_direct_vm_routing_allowlist_segment_boundary() -> None: assert _matches_direct_vm_prefix("playwright/execute", prefixes) is True assert _matches_direct_vm_prefix("process/exec", prefixes) is True assert _matches_direct_vm_prefix("process/proc-1/stdout/stream", prefixes) is True - assert _matches_direct_vm_prefix("fs/read", prefixes) is False + assert _matches_direct_vm_prefix("fs/read_file", prefixes) is True + assert _matches_direct_vm_prefix("fs/watch/watch-1/events", prefixes) is True + assert _matches_direct_vm_prefix("fsx/read_file", prefixes) is False + assert _matches_direct_vm_prefix("logs/stream", prefixes) is True + assert _matches_direct_vm_prefix("logs/stream/x", prefixes) is True + assert _matches_direct_vm_prefix("logs", prefixes) is False + assert _matches_direct_vm_prefix("logs/history", prefixes) is False + assert _matches_direct_vm_prefix("logstream", prefixes) is False + assert _matches_direct_vm_prefix("extensions", prefixes) is False + assert _matches_direct_vm_prefix("replays/rec-1", prefixes) is False def test_rewrite_direct_vm_options_keeps_telemetry_events_on_control_plane() -> None: @@ -435,7 +518,9 @@ def test_rewrite_direct_vm_options_keeps_telemetry_events_on_control_plane() -> cache = BrowserRouteCache() cache.set(BrowserRoute(session_id="sess-1", base_url="http://browser-session.test/browser/kernel", jwt="token-abc")) - config = BrowserRoutingConfig(subresources=("curl", "telemetry/stream", "computer", "playwright", "process")) + config = BrowserRoutingConfig( + subresources=("curl", "telemetry/stream", "computer", "playwright", "process", "fs", "logs/stream") + ) events = rewrite_direct_vm_options( FinalRequestOptions(method="get", url="/browsers/sess-1/telemetry/events"), cache=cache, config=config @@ -465,7 +550,32 @@ def test_rewrite_direct_vm_options_keeps_telemetry_events_on_control_plane() -> fs_read = rewrite_direct_vm_options( FinalRequestOptions(method="get", url="/browsers/sess-1/fs/read_file"), cache=cache, config=config ) - assert fs_read.url == "/browsers/sess-1/fs/read_file" + assert str(fs_read.url).startswith("http://browser-session.test/browser/kernel/fs/read_file") + + logs_stream = rewrite_direct_vm_options( + FinalRequestOptions(method="get", url="/browsers/sess-1/logs/stream"), cache=cache, config=config + ) + assert str(logs_stream.url).startswith("http://browser-session.test/browser/kernel/logs/stream") + + logs_root = rewrite_direct_vm_options( + FinalRequestOptions(method="get", url="/browsers/sess-1/logs"), cache=cache, config=config + ) + assert logs_root.url == "/browsers/sess-1/logs" + + logs_history = rewrite_direct_vm_options( + FinalRequestOptions(method="get", url="/browsers/sess-1/logs/history"), cache=cache, config=config + ) + assert logs_history.url == "/browsers/sess-1/logs/history" + + extensions = rewrite_direct_vm_options( + FinalRequestOptions(method="post", url="/browsers/sess-1/extensions"), cache=cache, config=config + ) + assert extensions.url == "/browsers/sess-1/extensions" + + replays = rewrite_direct_vm_options( + FinalRequestOptions(method="get", url="/browsers/sess-1/replays"), cache=cache, config=config + ) + assert replays.url == "/browsers/sess-1/replays" def test_browser_routing_config_from_env_empty_string_disables_routing(monkeypatch: pytest.MonkeyPatch) -> None: @@ -510,21 +620,24 @@ def test_default_browser_subresources_route_to_vm( @respx.mock -def test_fs_and_telemetry_events_stay_on_api_origin_by_default( +def test_control_plane_subresources_stay_on_api_origin_by_default( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) - fs_read = respx.get(f"{base_url}/browsers/sess-1/fs/read_file").mock( - return_value=httpx.Response(200, content=b"x", headers={"content-type": "application/octet-stream"}) - ) events = respx.get(f"{base_url}/browsers/sess-1/telemetry/events").mock(return_value=httpx.Response(200, json=[])) + replays = respx.get(f"{base_url}/browsers/sess-1/replays").mock(return_value=httpx.Response(200, json=[])) + extensions = respx.post(f"{base_url}/browsers/sess-1/extensions").mock(return_value=httpx.Response(204)) with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: _cache_browser(client) - client.browsers.fs.read_file("sess-1", path="/tmp/x") client.browsers.telemetry.events("sess-1") + client.browsers.replays.list("sess-1") + client.browsers.load_extensions("sess-1", extensions=[{"name": "ext", "zip_file": b"zip"}]) - assert fs_read.called assert events.called + assert replays.called + assert extensions.called + extensions_req = cast(httpx.Request, cast(Any, extensions.calls[0]).request) + assert extensions_req.headers.get("Authorization") == f"Bearer {api_key}" @respx.mock @@ -532,10 +645,6 @@ def test_stale_direct_vm_jwt_evicts_cache_and_retries_control_plane( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) - - def _skip_retry_sleep(_self: object, **_kwargs: object) -> None: - return None - monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep) vm = respx.post("http://browser-session.test/browser/kernel/computer/screenshot").mock( return_value=httpx.Response(401, text="Invalid JWT") @@ -596,3 +705,761 @@ def test_stale_direct_vm_auth_retry_does_not_require_cached_route() -> None: empty = BrowserRouteCache() assert should_retry_stale_direct_vm_auth(response) is True assert empty.get("sess-1") is None + + +@respx.mock +def test_fs_json_endpoints_route_to_vm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + list_files = respx.get("http://browser-session.test/browser/kernel/fs/list_files").mock( + return_value=httpx.Response(200, json=[]) + ) + move = respx.put("http://browser-session.test/browser/kernel/fs/move").mock(return_value=httpx.Response(204)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + client.browsers.fs.list_files("sess-1", path="/tmp") + client.browsers.fs.move("sess-1", dest_path="/tmp/b", src_path="/tmp/a") + + list_req = cast(httpx.Request, cast(Any, list_files.calls[0]).request) + assert list_req.url.params.get("path") == "/tmp" + assert list_req.url.params.get("jwt") == "token-abc" + assert list_req.headers.get("Authorization") is None + + move_req = cast(httpx.Request, cast(Any, move.calls[0]).request) + assert move_req.url.path == "/browser/kernel/fs/move" + assert move_req.url.params.get("jwt") == "token-abc" + assert move_req.headers.get("Authorization") is None + + +@respx.mock +def test_fs_read_file_routes_binary_response_from_vm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + read_file = respx.get("http://browser-session.test/browser/kernel/fs/read_file").mock( + return_value=httpx.Response(200, content=b"\x00binary", headers={"content-type": "application/octet-stream"}) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + response = client.browsers.fs.read_file("sess-1", path="/tmp/x") + + assert response.read() == b"\x00binary" + request = cast(httpx.Request, cast(Any, read_file.calls[0]).request) + assert request.url.params.get("path") == "/tmp/x" + assert request.url.params.get("jwt") == "token-abc" + assert request.headers.get("Authorization") is None + + +@respx.mock +def test_fs_write_file_routes_binary_body_to_vm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + write_file = respx.put("http://browser-session.test/browser/kernel/fs/write_file").mock( + return_value=httpx.Response(201) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + client.browsers.fs.write_file("sess-1", b"\x00payload", path="/tmp/x", mode="600") + + request = cast(httpx.Request, cast(Any, write_file.calls[0]).request) + assert request.content == b"\x00payload" + assert request.headers.get("content-type") == "application/octet-stream" + assert request.url.params.get("path") == "/tmp/x" + assert request.url.params.get("mode") == "600" + assert request.url.params.get("jwt") == "token-abc" + assert request.headers.get("Authorization") is None + + +@respx.mock +def test_fs_upload_routes_indexed_multipart_to_vm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + upload = respx.post("http://browser-session.test/browser/kernel/fs/upload").mock(return_value=httpx.Response(201)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + client.browsers.fs.upload( + "sess-1", + files=[ + {"dest_path": "/tmp/one", "file": b"one"}, + {"dest_path": "/tmp/two", "file": b"two"}, + ], + ) + + request = cast(httpx.Request, cast(Any, upload.calls[0]).request) + assert request.url.params.get("jwt") == "token-abc" + assert request.headers.get("Authorization") is None + body = request.read() + assert b'name="files[0][dest_path]"' in body + assert b'name="files[0][file]"' in body + assert b'name="files[1][dest_path]"' in body + assert b'name="files[1][file]"' in body + assert b"files[][" not in body + + +@respx.mock +def test_fs_watch_events_stream_routes_to_vm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + events = respx.get("http://browser-session.test/browser/kernel/fs/watch/watch-1/events").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"type":"CREATE","path":"/tmp/x","is_dir":false,"name":"x"}\n\n', + ) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + stream = client.browsers.fs.watch.events("watch-1", id_or_name="sess-1") + first = next(iter(stream)) + stream.close() + + assert first.path == "/tmp/x" + request = cast(httpx.Request, cast(Any, events.calls[0]).request) + assert request.url.path == "/browser/kernel/fs/watch/watch-1/events" + assert request.url.params.get("jwt") == "token-abc" + assert request.headers.get("Authorization") is None + + +@respx.mock +def test_logs_stream_routes_to_vm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + logs = respx.get("http://browser-session.test/browser/kernel/logs/stream").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"event":"log","message":"hello","timestamp":"2020-01-01T00:00:00Z"}\n\n', + ) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + stream = client.browsers.logs.stream("sess-1", source="path", path="/var/log/x", follow=True) + first = next(iter(stream)) + stream.close() + + assert first.message == "hello" + request = cast(httpx.Request, cast(Any, logs.calls[0]).request) + assert request.url.path == "/browser/kernel/logs/stream" + assert request.url.params.get("source") == "path" + assert request.url.params.get("path") == "/var/log/x" + assert request.url.params.get("follow") == "true" + assert request.url.params.get("jwt") == "token-abc" + assert request.headers.get("Authorization") is None + + +@pytest.mark.asyncio +async def test_async_logs_stream_cancellation_reaches_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + read_started = asyncio.Event() + transport_cancelled = asyncio.Event() + chunks: asyncio.Queue[bytes] = asyncio.Queue() + requested: list[httpx.URL] = [] + + class BlockingSSEStream(httpx.AsyncByteStream): + @override + async def __aiter__(self) -> AsyncIterator[bytes]: + read_started.set() + try: + while True: + yield await chunks.get() + except asyncio.CancelledError: + transport_cancelled.set() + raise + + @override + async def aclose(self) -> None: + pass + + async def handle_request(request: httpx.Request) -> httpx.Response: + requested.append(request.url) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=BlockingSSEStream(), + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handle_request)) + async with AsyncKernel( + base_url=base_url, + api_key=api_key, + http_client=http_client, + _strict_response_validation=True, + ) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + stream = await client.browsers.logs.stream("sess-1", source="supervisor", supervisor_process="chromium") + consumer = asyncio.create_task(stream.__anext__()) + await asyncio.wait_for(read_started.wait(), timeout=1) + + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(consumer, timeout=1) + await asyncio.wait_for(transport_cancelled.wait(), timeout=1) + + assert requested + assert requested[0].path == "/browser/kernel/logs/stream" + + +@respx.mock +def test_stale_direct_vm_jwt_replays_buffered_fs_body_on_control_plane( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep) + vm = respx.put("http://browser-session.test/browser/kernel/fs/write_file").mock( + return_value=httpx.Response(401, text="Invalid JWT") + ) + api = respx.put(f"{base_url}/browsers/sess-1/fs/write_file").mock(return_value=httpx.Response(201)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + client.browsers.fs.write_file("sess-1", b"payload", path="/tmp/x") + assert client.browser_route_cache.get("sess-1") is None + + assert vm.called + api_req = cast(httpx.Request, cast(Any, api.calls[0]).request) + assert api_req.content == b"payload" + assert api_req.url.params.get("path") == "/tmp/x" + assert api_req.url.params.get("jwt") is None + assert api_req.headers.get("Authorization") == f"Bearer {api_key}" + + +@respx.mock +def test_stale_direct_vm_jwt_does_not_replay_streamed_fs_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep) + vm = respx.put("http://browser-session.test/browser/kernel/fs/write_file").mock( + return_value=httpx.Response(401, text="Invalid JWT") + ) + api = respx.put(f"{base_url}/browsers/sess-1/fs/write_file").mock(return_value=httpx.Response(201)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + with pytest.raises(AuthenticationError): + client.browsers.fs.write_file("sess-1", iter([b"chunk-one", b"chunk-two"]), path="/tmp/x") + # The stale route is still evicted, so the caller's next attempt uses the control plane. + assert client.browser_route_cache.get("sess-1") is None + + assert vm.call_count == 1 + assert not api.called + + +@pytest.mark.asyncio +@respx.mock +async def test_async_stale_direct_vm_jwt_does_not_replay_streamed_fs_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + monkeypatch.setattr("kernel._base_client.AsyncAPIClient._sleep_for_retry", _skip_retry_sleep) + vm = respx.put("http://browser-session.test/browser/kernel/fs/write_file").mock( + return_value=httpx.Response(401, text="Invalid JWT") + ) + api = respx.put(f"{base_url}/browsers/sess-1/fs/write_file").mock(return_value=httpx.Response(201)) + + async def _chunks() -> AsyncIterator[bytes]: + yield b"chunk-one" + yield b"chunk-two" + + async with AsyncKernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + with pytest.raises(AuthenticationError): + await client.browsers.fs.write_file("sess-1", _chunks(), path="/tmp/x") + assert client.browser_route_cache.get("sess-1") is None + + assert vm.call_count == 1 + assert not api.called + + +@respx.mock +def test_stale_direct_vm_jwt_replays_multipart_upload_on_control_plane( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep) + vm = respx.post("http://browser-session.test/browser/kernel/fs/upload").mock( + return_value=httpx.Response(403, text="Invalid JWT") + ) + api = respx.post(f"{base_url}/browsers/sess-1/fs/upload").mock(return_value=httpx.Response(201)) + upload = tmp_path / "one.txt" + upload.write_bytes(b"file-bytes") + + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + with upload.open("rb") as handle: + client.browsers.fs.upload("sess-1", files=[{"dest_path": "/tmp/one", "file": handle}]) + + assert vm.called + api_req = cast(httpx.Request, cast(Any, api.calls[0]).request) + body = api_req.read() + assert b"file-bytes" in body + assert b'name="files[0][dest_path]"' in body + assert api_req.headers.get("Authorization") == f"Bearer {api_key}" + + +def test_direct_vm_request_body_is_replayable_classification(tmp_path: Path) -> None: + from kernel.lib.browser_routing.routing import direct_vm_request_body_is_replayable + + assert direct_vm_request_body_is_replayable(httpx.Request("GET", "http://vm.test/fs/read_file")) is True + assert direct_vm_request_body_is_replayable(httpx.Request("PUT", "http://vm.test/fs/write_file", content=b"x")) + assert ( + direct_vm_request_body_is_replayable(httpx.Request("PUT", "http://vm.test/fs/write_file", content=iter([b"x"]))) + is False + ) + + path = tmp_path / "one.txt" + path.write_bytes(b"file-bytes") + with path.open("rb") as handle: + seekable = httpx.Request( + "POST", + "http://vm.test/fs/upload", + data={"files[0][dest_path]": "/tmp/one"}, + files=[("files[0][file]", handle)], + ) + assert direct_vm_request_body_is_replayable(seekable) is True + + unseekable = httpx.Request( + "POST", + "http://vm.test/fs/upload", + files=[("files[0][file]", cast(Any, _UnseekableFile(b"file-bytes")))], + ) + assert direct_vm_request_body_is_replayable(unseekable) is False + + # seekable() is not proof: the rewind itself has to succeed. + lies_about_seekable = httpx.Request( + "POST", + "http://vm.test/fs/upload", + files=[("files[0][file]", cast(Any, _UnseekableFile(b"file-bytes", claims_seekable=True)))], + ) + assert direct_vm_request_body_is_replayable(lies_about_seekable) is False + + without_seekable_attr = httpx.Request( + "POST", + "http://vm.test/fs/upload", + files=[("files[0][file]", cast(Any, _NoSeekableAttrFile(b"file-bytes")))], + ) + assert direct_vm_request_body_is_replayable(without_seekable_attr) is False + + in_memory = io.BytesIO(b"file-bytes") + buffered = httpx.Request("POST", "http://vm.test/fs/upload", files=[("files[0][file]", in_memory)]) + assert direct_vm_request_body_is_replayable(buffered) is True + + in_memory.close() + assert direct_vm_request_body_is_replayable(buffered) is False + + +@respx.mock +def test_env_override_can_exclude_fs_and_logs(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", "computer") + fs_read = respx.get(f"{base_url}/browsers/sess-1/fs/read_file").mock( + return_value=httpx.Response(200, content=b"x", headers={"content-type": "application/octet-stream"}) + ) + logs = respx.get(f"{base_url}/browsers/sess-1/logs/stream").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"event":"log","message":"hello","timestamp":"2020-01-01T00:00:00Z"}\n\n', + ) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + client.browsers.fs.read_file("sess-1", path="/tmp/x") + client.browsers.logs.stream("sess-1", source="path", path="/var/log/x").close() + + assert fs_read.called + assert logs.called + + +@respx.mock +def test_empty_env_disables_fs_and_logs_routing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", "") + fs_read = respx.get(f"{base_url}/browsers/sess-1/fs/read_file").mock( + return_value=httpx.Response(200, content=b"x", headers={"content-type": "application/octet-stream"}) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + response = client.browsers.fs.read_file("sess-1", path="/tmp/x") + + assert fs_read.called + request = cast(httpx.Request, cast(Any, fs_read.calls[0]).request) + assert request.url.params.get("jwt") is None + assert request.headers.get("Authorization") == f"Bearer {api_key}" + assert response.read() == b"x" + + +@respx.mock +def test_stale_direct_vm_jwt_does_not_replay_multipart_that_cannot_rewind( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep) + vm = respx.post("http://browser-session.test/browser/kernel/fs/upload").mock( + return_value=httpx.Response(401, text="Invalid JWT") + ) + api = respx.post(f"{base_url}/browsers/sess-1/fs/upload").mock(return_value=httpx.Response(201)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + with pytest.raises(AuthenticationError): + client.browsers.fs.upload( + "sess-1", + files=[ + { + "dest_path": "/tmp/one", + "file": cast(Any, _UnseekableFile(b"file-bytes", claims_seekable=True)), + } + ], + ) + assert client.browser_route_cache.get("sess-1") is None + + assert vm.call_count == 1 + assert not api.called + + +@pytest.mark.asyncio +@respx.mock +async def test_async_stale_direct_vm_jwt_does_not_replay_multipart_that_cannot_rewind( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + monkeypatch.setattr("kernel._base_client.AsyncAPIClient._sleep_for_retry", _skip_retry_sleep) + vm = respx.post("http://browser-session.test/browser/kernel/fs/upload").mock( + return_value=httpx.Response(403, text="Invalid JWT") + ) + api = respx.post(f"{base_url}/browsers/sess-1/fs/upload").mock(return_value=httpx.Response(201)) + async with AsyncKernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + with pytest.raises(PermissionDeniedError): + await client.browsers.fs.upload( + "sess-1", + files=[ + { + "dest_path": "/tmp/one", + "file": cast(Any, _UnseekableFile(b"file-bytes", claims_seekable=True)), + } + ], + ) + assert client.browser_route_cache.get("sess-1") is None + + assert vm.call_count == 1 + assert not api.called + + +@respx.mock +def test_stale_direct_vm_jwt_evicts_route_without_retries_for_buffered_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + vm = respx.put("http://browser-session.test/browser/kernel/fs/write_file").mock( + return_value=httpx.Response(401, text="Invalid JWT") + ) + api = respx.put(f"{base_url}/browsers/sess-1/fs/write_file").mock(return_value=httpx.Response(201)) + with Kernel(base_url=base_url, api_key=api_key, max_retries=0, _strict_response_validation=True) as client: + _cache_browser(client) + with pytest.raises(AuthenticationError): + client.browsers.fs.write_file("sess-1", b"payload", path="/tmp/x") + assert client.browser_route_cache.get("sess-1") is None + + # The caller's next attempt goes to the control plane. + client.browsers.fs.write_file("sess-1", b"payload", path="/tmp/x") + + assert vm.call_count == 1 + api_req = cast(httpx.Request, cast(Any, api.calls[0]).request) + assert api_req.content == b"payload" + assert api_req.url.params.get("jwt") is None + assert api_req.headers.get("Authorization") == f"Bearer {api_key}" + + +@pytest.mark.asyncio +@respx.mock +async def test_async_stale_direct_vm_jwt_evicts_route_without_retries_for_streamed_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + vm = respx.put("http://browser-session.test/browser/kernel/fs/write_file").mock( + return_value=httpx.Response(403, text="Invalid JWT") + ) + api = respx.put(f"{base_url}/browsers/sess-1/fs/write_file").mock(return_value=httpx.Response(201)) + + async def _chunks() -> AsyncIterator[bytes]: + yield b"chunk-one" + + async with AsyncKernel( + base_url=base_url, api_key=api_key, max_retries=0, _strict_response_validation=True + ) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + with pytest.raises(PermissionDeniedError): + await client.browsers.fs.write_file("sess-1", _chunks(), path="/tmp/x") + assert client.browser_route_cache.get("sess-1") is None + + await client.browsers.fs.write_file("sess-1", b"payload", path="/tmp/x") + + assert vm.call_count == 1 + api_req = cast(httpx.Request, cast(Any, api.calls[0]).request) + assert api_req.url.params.get("jwt") is None + assert api_req.headers.get("Authorization") == f"Bearer {api_key}" + + +@respx.mock +def test_load_extensions_multipart_encoding_is_unchanged(monkeypatch: pytest.MonkeyPatch) -> None: + # Indexed names are scoped to fs.upload; every other multipart endpoint keeps + # the client's generic array encoding. + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + extensions = respx.post(f"{base_url}/browsers/sess-1/extensions").mock(return_value=httpx.Response(204)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + _cache_browser(client) + client.browsers.load_extensions( + "sess-1", + extensions=[ + {"name": "one", "zip_file": b"zip-one"}, + {"name": "two", "zip_file": b"zip-two"}, + ], + ) + + body = cast(httpx.Request, cast(Any, extensions.calls[0]).request).read() + assert b'name="extensions[][name]"' in body + assert b'name="extensions[][zip_file]"' in body + assert b"extensions[0]" not in body + + +def test_generic_multipart_array_encoding_is_unchanged() -> None: + from kernel._models import FinalRequestOptions + + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + request = client._build_request( # pyright: ignore[reportPrivateUsage] + FinalRequestOptions.construct( + method="post", + url="/foo", + headers={"Content-Type": "multipart/form-data; boundary=abc"}, + json_data={"array": ["foo", "bar"]}, + files=[("foo.txt", b"hello world")], + ) + ) + + body = request.read() + assert b'name="array[]"' in body + assert b'name="array[0]"' not in body + + +def test_indexed_multipart_body_flattens_only_given_values() -> None: + from kernel._types import omit + from kernel.lib.multipart import indexed_multipart_body + + assert indexed_multipart_body( + { + "files": [ + {"dest_path": "/tmp/one", "mode": omit}, + {"dest_path": "/tmp/two"}, + ], + "flag": True, + "skipped": omit, + } + ) == { + "files[0][dest_path]": "/tmp/one", + "files[1][dest_path]": "/tmp/two", + "flag": True, + } + + +class _FailingSyncStream(httpx.SyncByteStream): + """A response body that fails while it is being read.""" + + @override + def __iter__(self) -> Iterator[bytes]: + raise httpx.ReadError("connection reset while reading the error body") + + +class _FailingAsyncStream(httpx.AsyncByteStream): + @override + async def __aiter__(self) -> AsyncIterator[bytes]: + raise httpx.ReadError("connection reset while reading the error body") + yield b"" # pragma: no cover - unreachable, keeps this an async generator + + +def test_stale_direct_vm_jwt_evicts_route_when_error_body_read_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + requests: list[httpx.Request] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append(request) + if "browser-session.test" in str(request.url): + return httpx.Response(401, stream=_FailingSyncStream(), headers={"content-type": "text/plain"}) + return httpx.Response(200, content=b"png", headers={"content-type": "image/png"}) + + http_client = httpx.Client(transport=httpx.MockTransport(handle_request)) + with Kernel( + base_url=base_url, + api_key=api_key, + max_retries=0, + http_client=http_client, + _strict_response_validation=True, + ) as client: + _cache_browser(client) + with pytest.raises(APIConnectionError): + client.browsers.computer.capture_screenshot("sess-1") + # The status was known before the body read failed, so the dead route is gone. + assert client.browser_route_cache.get("sess-1") is None + + client.browsers.computer.capture_screenshot("sess-1") + + assert str(requests[0].url).startswith("http://browser-session.test/browser/kernel/computer/screenshot") + assert requests[1].url == httpx.URL(f"{base_url}/browsers/sess-1/computer/screenshot") + assert requests[1].url.params.get("jwt") is None + assert requests[1].headers.get("Authorization") == f"Bearer {api_key}" + + +@pytest.mark.asyncio +async def test_async_stale_direct_vm_jwt_evicts_route_when_error_body_read_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + requests: list[httpx.Request] = [] + + async def handle_request(request: httpx.Request) -> httpx.Response: + requests.append(request) + if "browser-session.test" in str(request.url): + return httpx.Response(403, stream=_FailingAsyncStream(), headers={"content-type": "text/plain"}) + return httpx.Response(200, content=b"png", headers={"content-type": "image/png"}) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handle_request)) + async with AsyncKernel( + base_url=base_url, + api_key=api_key, + max_retries=0, + http_client=http_client, + _strict_response_validation=True, + ) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + with pytest.raises(APIConnectionError): + await client.browsers.computer.capture_screenshot("sess-1") + assert client.browser_route_cache.get("sess-1") is None + + await client.browsers.computer.capture_screenshot("sess-1") + + assert str(requests[0].url).startswith("http://browser-session.test/browser/kernel/computer/screenshot") + assert requests[1].url == httpx.URL(f"{base_url}/browsers/sess-1/computer/screenshot") + assert requests[1].url.params.get("jwt") is None + assert requests[1].headers.get("Authorization") == f"Bearer {api_key}" + + +def test_copied_client_registers_one_route_eviction_hook() -> None: + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + copied = client.copy(api_key="sk-456") + assert copied.browser_route_cache is client.browser_route_cache + hooks = client._client.event_hooks["response"] # pyright: ignore[reportPrivateUsage] + assert len(hooks) == 1 + + +def test_route_eviction_hook_runs_before_caller_response_hooks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + requests: list[httpx.Request] = [] + caller_hook_statuses: list[int] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append(request) + if "browser-session.test" in str(request.url): + return httpx.Response(401, stream=_FailingSyncStream(), headers={"content-type": "text/plain"}) + return httpx.Response(200, content=b"png", headers={"content-type": "image/png"}) + + def caller_hook(response: httpx.Response) -> None: + caller_hook_statuses.append(response.status_code) + if response.status_code in {401, 403}: + # Reading a failing body raises out of the hook chain, which would + # skip any eviction hook registered after this one. + response.read() + + http_client = httpx.Client( + transport=httpx.MockTransport(handle_request), + event_hooks={"response": [caller_hook]}, + ) + with Kernel( + base_url=base_url, + api_key=api_key, + max_retries=0, + http_client=http_client, + _strict_response_validation=True, + ) as client: + _cache_browser(client) + assert http_client.event_hooks["response"][-1] is caller_hook + with pytest.raises(APIConnectionError): + client.browsers.computer.capture_screenshot("sess-1") + assert caller_hook_statuses == [401] + assert client.browser_route_cache.get("sess-1") is None + + client.browsers.computer.capture_screenshot("sess-1") + + assert str(requests[0].url).startswith("http://browser-session.test/browser/kernel/computer/screenshot") + assert requests[1].url == httpx.URL(f"{base_url}/browsers/sess-1/computer/screenshot") + assert requests[1].url.params.get("jwt") is None + assert requests[1].headers.get("Authorization") == f"Bearer {api_key}" + + +@pytest.mark.asyncio +async def test_async_route_eviction_hook_runs_before_caller_response_hooks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False) + requests: list[httpx.Request] = [] + caller_hook_statuses: list[int] = [] + + async def handle_request(request: httpx.Request) -> httpx.Response: + requests.append(request) + if "browser-session.test" in str(request.url): + return httpx.Response(403, stream=_FailingAsyncStream(), headers={"content-type": "text/plain"}) + return httpx.Response(200, content=b"png", headers={"content-type": "image/png"}) + + async def caller_hook(response: httpx.Response) -> None: + caller_hook_statuses.append(response.status_code) + if response.status_code in {401, 403}: + await response.aread() + + http_client = httpx.AsyncClient( + transport=httpx.MockTransport(handle_request), + event_hooks={"response": [caller_hook]}, + ) + async with AsyncKernel( + base_url=base_url, + api_key=api_key, + max_retries=0, + http_client=http_client, + _strict_response_validation=True, + ) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + assert http_client.event_hooks["response"][-1] is caller_hook + with pytest.raises(APIConnectionError): + await client.browsers.computer.capture_screenshot("sess-1") + assert caller_hook_statuses == [403] + assert client.browser_route_cache.get("sess-1") is None + + await client.browsers.computer.capture_screenshot("sess-1") + + assert str(requests[0].url).startswith("http://browser-session.test/browser/kernel/computer/screenshot") + assert requests[1].url == httpx.URL(f"{base_url}/browsers/sess-1/computer/screenshot") + assert requests[1].url.params.get("jwt") is None + assert requests[1].headers.get("Authorization") == f"Bearer {api_key}" + + +def test_route_eviction_hook_is_registered_once_before_caller_hooks() -> None: + def caller_hook(_response: httpx.Response) -> None: # pragma: no cover - never invoked + return None + + http_client = httpx.Client(event_hooks={"response": [caller_hook]}) + with Kernel( + base_url=base_url, + api_key=api_key, + http_client=http_client, + _strict_response_validation=True, + ) as client: + client.copy(api_key="sk-456") + hooks = http_client.event_hooks["response"] + assert len(hooks) == 2 + assert hooks[1] is caller_hook