From 9088723937cde444617ed97d3188e78c705e5cdd Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Wed, 2 Sep 2026 06:17:24 +0000 Subject: [PATCH 1/5] Route browser fs and logs endpoints directly to the VM Add fs and logs to the default direct-to-VM subresource prefixes so filesystem operations and log streaming use the cached browser base_url and JWT instead of the control plane. Serialize multipart array entries with indexed names so a file part stays associated with the sibling fields of its array entry, which repeated `files[][file]` names cannot express. Only retry a stale direct-to-VM auth failure on the control plane when the request body can be rebuilt byte for byte; a streamed body is consumed by the direct attempt, so retrying would send a truncated body. The stale route is evicted either way. --- src/kernel/_base_client.py | 5 +- src/kernel/_client.py | 15 +- src/kernel/_utils/_utils.py | 6 +- src/kernel/lib/browser_routing/routing.py | 56 ++- tests/test_browser_routing.py | 435 +++++++++++++++++++++- tests/test_client.py | 16 +- tests/test_extract_files.py | 6 +- 7 files changed, 504 insertions(+), 35 deletions(-) diff --git a/src/kernel/_base_client.py b/src/kernel/_base_client.py index 2599dc41..50b190a5 100644 --- a/src/kernel/_base_client.py +++ b/src/kernel/_base_client.py @@ -587,7 +587,10 @@ def _serialize_multipartform(self, data: Mapping[object, object]) -> dict[str, o # TODO: type ignore is required as stringify_items is well typed but we can't be # well typed without heavy validation. data, # type: ignore - array_format="brackets", + # Indexed names (`files[0][dest_path]`) keep each array entry's fields + # grouped together; repeated `files[][dest_path]` parts cannot be + # matched back to their file part. `extract_files` uses the same format. + array_format="indices", ) serialized: dict[str, object] = {} for key, value in items: diff --git a/src/kernel/_client.py b/src/kernel/_client.py index 2c06cb04..995e2dad 100644 --- a/src/kernel/_client.py +++ b/src/kernel/_client.py @@ -41,6 +41,7 @@ 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, maybe_evict_browser_route_from_response, maybe_populate_browser_route_cache_from_response, @@ -365,9 +366,12 @@ 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): + if is_stale_direct_vm_auth_response(response): maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) - return True + # The route is evicted either way; only retry 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 @@ -748,9 +752,12 @@ 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): + if is_stale_direct_vm_auth_response(response): maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) - return True + # The route is evicted either way; only retry 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/_utils/_utils.py b/src/kernel/_utils/_utils.py index 199cd231..45116559 100644 --- a/src/kernel/_utils/_utils.py +++ b/src/kernel/_utils/_utils.py @@ -40,7 +40,7 @@ def extract_files( query: Mapping[str, object], *, paths: Sequence[Sequence[str]], - array_format: ArrayFormat = "brackets", + array_format: ArrayFormat = "indices", ) -> list[tuple[str, FileTypes]]: """Recursively extract files from the given dictionary based on specified paths. @@ -48,7 +48,9 @@ def extract_files( ``array_format`` controls how ```` segments contribute to the emitted field name. Supported values: ``"brackets"`` (``foo[]``), ``"repeat"`` and - ``"comma"`` (``foo``), ``"indices"`` (``foo[0]``, ``foo[1]``). + ``"comma"`` (``foo``), ``"indices"`` (``foo[0]``, ``foo[1]``). Indexed names are + the default so that a file part stays associated with the sibling fields of the + same array entry, which repeated ``foo[]`` names cannot express. Note: this mutates the given dictionary. """ diff --git a/src/kernel/lib/browser_routing/routing.py b/src/kernel/lib/browser_routing/routing.py index bad5e4ea..1fa69f54 100644 --- a/src/kernel/lib/browser_routing/routing.py +++ b/src/kernel/lib/browser_routing/routing.py @@ -44,7 +44,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", + ) + ) if raw.strip() == "": return BrowserRoutingConfig() @@ -189,7 +199,49 @@ def is_stale_direct_vm_auth_response(response: httpx.Response) -> bool: 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 + if not callable(getattr(file, "seek", None)): + return False + seekable = getattr(file, "seekable", None) + # httpx rewinds seekable file fields before rendering them again. + return bool(seekable()) if callable(seekable) else True def _session_id_from_browser_delete_path(path: str) -> str | None: diff --git a/tests/test_browser_routing.py b/tests/test_browser_routing.py index eb47e040..4e154802 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -3,13 +3,14 @@ import os import asyncio from typing import Any, 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, AuthenticationError, InternalServerError from kernel.lib.browser_routing.util import jwt_from_cdp_ws_url from kernel.lib.browser_routing.routing import ( BrowserRoute, @@ -37,6 +38,30 @@ def _fake_browser() -> dict[str, object]: } +def _skip_retry_sleep(_self: object, **_kwargs: object) -> None: + return None + + +class _UnseekableFile: + """A file-like upload body that cannot be rewound, e.g. a pipe.""" + + name = "one.txt" + + def __init__(self, content: bytes) -> None: + self._content = content + + 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 + + def seekable(self) -> bool: + return False + + 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 +423,8 @@ def test_browser_routing_config_from_env_defaults(monkeypatch: pytest.MonkeyPatc "computer", "playwright", "process", + "fs", + "logs", ) @@ -407,7 +434,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") 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 +445,13 @@ 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("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 +468,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") + ) events = rewrite_direct_vm_options( FinalRequestOptions(method="get", url="/browsers/sess-1/telemetry/events"), cache=cache, config=config @@ -465,7 +500,22 @@ 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") + + 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 +560,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 +585,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 +645,359 @@ 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(204) + ) + 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(204)) + 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(204)) + 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(204)) + 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(204)) + + 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(204)) + 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 + + +@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" diff --git a/tests/test_client.py b/tests/test_client.py index c3a1186c..3d5f915b 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -561,11 +561,11 @@ def test_multipart_repeating_array(self, client: Kernel) -> None: assert request.read().split(b"\r\n") == [ b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[]"', + b'Content-Disposition: form-data; name="array[0]"', b"", b"foo", b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[]"', + b'Content-Disposition: form-data; name="array[1]"', b"", b"bar", b"--6b7ba517decee4a450543ea6ae821c82", @@ -880,7 +880,7 @@ def test_parse_retry_after_header( calculated = client._calculate_retry_timeout(remaining_retries, options, headers) assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType] - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: Kernel) -> None: @@ -891,7 +891,7 @@ def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, clien assert _get_open_connections(client) == 0 - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: Kernel) -> None: @@ -1488,11 +1488,11 @@ def test_multipart_repeating_array(self, async_client: AsyncKernel) -> None: assert request.read().split(b"\r\n") == [ b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[]"', + b'Content-Disposition: form-data; name="array[0]"', b"", b"foo", b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[]"', + b'Content-Disposition: form-data; name="array[1]"', b"", b"bar", b"--6b7ba517decee4a450543ea6ae821c82", @@ -1824,7 +1824,7 @@ async def test_parse_retry_after_header( calculated = async_client._calculate_retry_timeout(remaining_retries, options, headers) assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType] - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncKernel) -> None: @@ -1835,7 +1835,7 @@ async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, assert _get_open_connections(async_client) == 0 - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncKernel) -> None: diff --git a/tests/test_extract_files.py b/tests/test_extract_files.py index 54ef03af..1325ec06 100644 --- a/tests/test_extract_files.py +++ b/tests/test_extract_files.py @@ -29,15 +29,15 @@ def test_removes_files_from_input() -> None: def test_multiple_files() -> None: query = {"documents": [{"file": b"My first file"}, {"file": b"My second file"}]} assert extract_files(query, paths=[["documents", "", "file"]]) == [ - ("documents[][file]", b"My first file"), - ("documents[][file]", b"My second file"), + ("documents[0][file]", b"My first file"), + ("documents[1][file]", b"My second file"), ] assert query == {"documents": [{}, {}]} def test_top_level_file_array() -> None: query = {"files": [b"file one", b"file two"], "title": "hello"} - assert extract_files(query, paths=[["files", ""]]) == [("files[]", b"file one"), ("files[]", b"file two")] + assert extract_files(query, paths=[["files", ""]]) == [("files[0]", b"file one"), ("files[1]", b"file two")] assert query == {"title": "hello"} From 96006001897a6f33da15b77eda0095fe0bb50ff3 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Wed, 2 Sep 2026 06:51:37 +0000 Subject: [PATCH 2/5] Address review: narrow multipart change, prove rewinds, always evict Scope the indexed multipart array names to fs.upload instead of changing the client's generic array encoding: the endpoint now flattens its own body with indexed names and asks extract_files for matching file part names, so load_extensions and any other multipart array keep their existing wire format. Prove a multipart file field can be rewound before treating a stale-JWT failure as retryable. A wrapper can report seekable() while seek() raises, which rendered the fallback body as an empty part. Evict a stale direct-to-VM route from the terminal error path too. Retry eligibility is only consulted when retries remain, so with max_retries=0 a VM 401/403 previously left the dead route cached and wedged later calls. Route only logs/stream rather than the whole logs subresource. --- src/kernel/_base_client.py | 5 +- src/kernel/_client.py | 28 ++- src/kernel/_utils/_utils.py | 6 +- src/kernel/lib/browser_routing/routing.py | 39 ++- src/kernel/lib/multipart.py | 41 ++++ src/kernel/resources/browsers/fs/fs.py | 19 +- tests/test_browser_routing.py | 274 +++++++++++++++++++++- tests/test_client.py | 16 +- tests/test_extract_files.py | 6 +- 9 files changed, 396 insertions(+), 38 deletions(-) create mode 100644 src/kernel/lib/multipart.py diff --git a/src/kernel/_base_client.py b/src/kernel/_base_client.py index 50b190a5..2599dc41 100644 --- a/src/kernel/_base_client.py +++ b/src/kernel/_base_client.py @@ -587,10 +587,7 @@ def _serialize_multipartform(self, data: Mapping[object, object]) -> dict[str, o # TODO: type ignore is required as stringify_items is well typed but we can't be # well typed without heavy validation. data, # type: ignore - # Indexed names (`files[0][dest_path]`) keep each array entry's fields - # grouped together; repeated `files[][dest_path]` parts cannot be - # matched back to their file part. `extract_files` uses the same format. - array_format="indices", + array_format="brackets", ) serialized: dict[str, object] = {} for key, value in items: diff --git a/src/kernel/_client.py b/src/kernel/_client.py index 995e2dad..07fd085b 100644 --- a/src/kernel/_client.py +++ b/src/kernel/_client.py @@ -364,16 +364,28 @@ def _prepare_options(self, options: Any) -> Any: def _prepare_request(self, request: httpx.Request) -> None: strip_direct_vm_auth(request, cache=self.browser_route_cache) + def _evict_stale_direct_vm_route(self, response: httpx.Response) -> None: + maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) + @override def _should_retry(self, response: httpx.Response) -> bool: if is_stale_direct_vm_auth_response(response): - maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) + self._evict_stale_direct_vm_route(response) # The route is evicted either way; only retry 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 + def _make_status_error_from_response(self, response: httpx.Response) -> APIStatusError: + # `_should_retry` never runs when the request has no retries left, so this + # is the only place a stale direct-to-VM route gets evicted before the + # error surfaces to the caller. + if is_stale_direct_vm_auth_response(response): + self._evict_stale_direct_vm_route(response) + return super()._make_status_error_from_response(response) + @override def _process_response( self, @@ -750,16 +762,28 @@ async def _prepare_options(self, options: Any) -> Any: async def _prepare_request(self, request: httpx.Request) -> None: strip_direct_vm_auth(request, cache=self.browser_route_cache) + def _evict_stale_direct_vm_route(self, response: httpx.Response) -> None: + maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) + @override def _should_retry(self, response: httpx.Response) -> bool: if is_stale_direct_vm_auth_response(response): - maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) + self._evict_stale_direct_vm_route(response) # The route is evicted either way; only retry 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 + def _make_status_error_from_response(self, response: httpx.Response) -> APIStatusError: + # `_should_retry` never runs when the request has no retries left, so this + # is the only place a stale direct-to-VM route gets evicted before the + # error surfaces to the caller. + if is_stale_direct_vm_auth_response(response): + self._evict_stale_direct_vm_route(response) + return super()._make_status_error_from_response(response) + @override async def _process_response( self, diff --git a/src/kernel/_utils/_utils.py b/src/kernel/_utils/_utils.py index 45116559..199cd231 100644 --- a/src/kernel/_utils/_utils.py +++ b/src/kernel/_utils/_utils.py @@ -40,7 +40,7 @@ def extract_files( query: Mapping[str, object], *, paths: Sequence[Sequence[str]], - array_format: ArrayFormat = "indices", + array_format: ArrayFormat = "brackets", ) -> list[tuple[str, FileTypes]]: """Recursively extract files from the given dictionary based on specified paths. @@ -48,9 +48,7 @@ def extract_files( ``array_format`` controls how ```` segments contribute to the emitted field name. Supported values: ``"brackets"`` (``foo[]``), ``"repeat"`` and - ``"comma"`` (``foo``), ``"indices"`` (``foo[0]``, ``foo[1]``). Indexed names are - the default so that a file part stays associated with the sibling fields of the - same array entry, which repeated ``foo[]`` names cannot express. + ``"comma"`` (``foo``), ``"indices"`` (``foo[0]``, ``foo[1]``). Note: this mutates the given dictionary. """ diff --git a/src/kernel/lib/browser_routing/routing.py b/src/kernel/lib/browser_routing/routing.py index 1fa69f54..f005644f 100644 --- a/src/kernel/lib/browser_routing/routing.py +++ b/src/kernel/lib/browser_routing/routing.py @@ -52,7 +52,7 @@ def browser_routing_config_from_env() -> BrowserRoutingConfig: "playwright", "process", "fs", - "logs", + "logs/stream", ) ) if raw.strip() == "": @@ -237,11 +237,40 @@ def _multipart_field_is_replayable(field: Any) -> bool: return True if getattr(file, "closed", False): return False - if not callable(getattr(file, "seek", None)): + 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 - seekable = getattr(file, "seekable", None) - # httpx rewinds seekable file fields before rendering them again. - return bool(seekable()) if callable(seekable) else True + + 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 4e154802..756fbcda 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -1,5 +1,6 @@ from __future__ import annotations +import io import os import asyncio from typing import Any, AsyncIterator, cast @@ -10,7 +11,13 @@ import respx import pytest -from kernel import Kernel, AsyncKernel, AuthenticationError, InternalServerError +from kernel import ( + Kernel, + AsyncKernel, + AuthenticationError, + InternalServerError, + PermissionDeniedError, +) from kernel.lib.browser_routing.util import jwt_from_cdp_ws_url from kernel.lib.browser_routing.routing import ( BrowserRoute, @@ -42,22 +49,61 @@ def _skip_retry_sleep(_self: object, **_kwargs: object) -> None: return None -class _UnseekableFile: - """A file-like upload body that cannot be rewound, e.g. a pipe.""" +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) -> None: + 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 False + 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") @@ -424,7 +470,7 @@ def test_browser_routing_config_from_env_defaults(monkeypatch: pytest.MonkeyPatc "playwright", "process", "fs", - "logs", + "logs/stream", ) @@ -434,7 +480,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", "fs", "logs") + 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 @@ -449,6 +495,9 @@ def test_direct_vm_routing_allowlist_segment_boundary() -> None: 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 @@ -469,7 +518,7 @@ 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", "fs", "logs") + subresources=("curl", "telemetry/stream", "computer", "playwright", "process", "fs", "logs/stream") ) events = rewrite_direct_vm_options( @@ -507,6 +556,16 @@ def test_rewrite_direct_vm_options_keeps_telemetry_events_on_control_plane() -> ) 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 ) @@ -963,6 +1022,28 @@ def test_direct_vm_request_body_is_replayable_classification(tmp_path: Path) -> ) 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: @@ -1001,3 +1082,180 @@ def test_empty_env_disables_fs_and_logs_routing(monkeypatch: pytest.MonkeyPatch) 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(204)) + 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(204)) + 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(204)) + 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(204)) + + 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, + } diff --git a/tests/test_client.py b/tests/test_client.py index 3d5f915b..c3a1186c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -561,11 +561,11 @@ def test_multipart_repeating_array(self, client: Kernel) -> None: assert request.read().split(b"\r\n") == [ b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[0]"', + b'Content-Disposition: form-data; name="array[]"', b"", b"foo", b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[1]"', + b'Content-Disposition: form-data; name="array[]"', b"", b"bar", b"--6b7ba517decee4a450543ea6ae821c82", @@ -880,7 +880,7 @@ def test_parse_retry_after_header( calculated = client._calculate_retry_timeout(remaining_retries, options, headers) assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType] - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: Kernel) -> None: @@ -891,7 +891,7 @@ def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, clien assert _get_open_connections(client) == 0 - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: Kernel) -> None: @@ -1488,11 +1488,11 @@ def test_multipart_repeating_array(self, async_client: AsyncKernel) -> None: assert request.read().split(b"\r\n") == [ b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[0]"', + b'Content-Disposition: form-data; name="array[]"', b"", b"foo", b"--6b7ba517decee4a450543ea6ae821c82", - b'Content-Disposition: form-data; name="array[1]"', + b'Content-Disposition: form-data; name="array[]"', b"", b"bar", b"--6b7ba517decee4a450543ea6ae821c82", @@ -1824,7 +1824,7 @@ async def test_parse_retry_after_header( calculated = async_client._calculate_retry_timeout(remaining_retries, options, headers) assert calculated == pytest.approx(timeout, 0.5 * 0.875) # pyright: ignore[reportUnknownMemberType] - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncKernel) -> None: @@ -1835,7 +1835,7 @@ async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, assert _get_open_connections(async_client) == 0 - @pytest.mark.skip() # SDK-2615 + @pytest.mark.skip() # SDK-2615 @mock.patch("kernel._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, async_client: AsyncKernel) -> None: diff --git a/tests/test_extract_files.py b/tests/test_extract_files.py index 1325ec06..54ef03af 100644 --- a/tests/test_extract_files.py +++ b/tests/test_extract_files.py @@ -29,15 +29,15 @@ def test_removes_files_from_input() -> None: def test_multiple_files() -> None: query = {"documents": [{"file": b"My first file"}, {"file": b"My second file"}]} assert extract_files(query, paths=[["documents", "", "file"]]) == [ - ("documents[0][file]", b"My first file"), - ("documents[1][file]", b"My second file"), + ("documents[][file]", b"My first file"), + ("documents[][file]", b"My second file"), ] assert query == {"documents": [{}, {}]} def test_top_level_file_array() -> None: query = {"files": [b"file one", b"file two"], "title": "hello"} - assert extract_files(query, paths=[["files", ""]]) == [("files[0]", b"file one"), ("files[1]", b"file two")] + assert extract_files(query, paths=[["files", ""]]) == [("files[]", b"file one"), ("files[]", b"file two")] assert query == {"title": "hello"} From c01804d33d22237b6912d4479fddcfbed2c3cd72 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Wed, 2 Sep 2026 07:08:10 +0000 Subject: [PATCH 3/5] Evict a stale direct-to-VM route before the error body is read httpx reads the body of a non-streamed response inside send(), so a 401 or 403 whose body read fails surfaces as a connection error and never reaches the status-error path. The route eviction ran after that read, which left a dead JWT cached and wedged every later call for the session. Move eviction into an httpx response event hook, which runs once the status is known and before any body is read, for both the sync and async clients. `_should_retry` now only decides whether replaying the body on the control plane is safe. --- src/kernel/_client.py | 42 +++------- src/kernel/lib/browser_routing/routing.py | 45 +++++++++++ tests/test_browser_routing.py | 96 ++++++++++++++++++++++- 3 files changed, 150 insertions(+), 33 deletions(-) diff --git a/src/kernel/_client.py b/src/kernel/_client.py index 07fd085b..dc506418 100644 --- a/src/kernel/_client.py +++ b/src/kernel/_client.py @@ -43,7 +43,9 @@ 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, ) @@ -204,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: @@ -364,28 +367,15 @@ def _prepare_options(self, options: Any) -> Any: def _prepare_request(self, request: httpx.Request) -> None: strip_direct_vm_auth(request, cache=self.browser_route_cache) - def _evict_stale_direct_vm_route(self, response: httpx.Response) -> None: - maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) - @override def _should_retry(self, response: httpx.Response) -> bool: if is_stale_direct_vm_auth_response(response): - self._evict_stale_direct_vm_route(response) - # The route is evicted either way; only retry when the body can be - # rebuilt, otherwise the caller sees the original auth failure and a - # later call goes to the control plane. + # 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 - def _make_status_error_from_response(self, response: httpx.Response) -> APIStatusError: - # `_should_retry` never runs when the request has no retries left, so this - # is the only place a stale direct-to-VM route gets evicted before the - # error surfaces to the caller. - if is_stale_direct_vm_auth_response(response): - self._evict_stale_direct_vm_route(response) - return super()._make_status_error_from_response(response) - @override def _process_response( self, @@ -602,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: @@ -762,28 +753,15 @@ async def _prepare_options(self, options: Any) -> Any: async def _prepare_request(self, request: httpx.Request) -> None: strip_direct_vm_auth(request, cache=self.browser_route_cache) - def _evict_stale_direct_vm_route(self, response: httpx.Response) -> None: - maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache) - @override def _should_retry(self, response: httpx.Response) -> bool: if is_stale_direct_vm_auth_response(response): - self._evict_stale_direct_vm_route(response) - # The route is evicted either way; only retry when the body can be - # rebuilt, otherwise the caller sees the original auth failure and a - # later call goes to the control plane. + # 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 - def _make_status_error_from_response(self, response: httpx.Response) -> APIStatusError: - # `_should_retry` never runs when the request has no retries left, so this - # is the only place a stale direct-to-VM route gets evicted before the - # error surfaces to the caller. - if is_stale_direct_vm_auth_response(response): - self._evict_stale_direct_vm_route(response) - return super()._make_status_error_from_response(response) - @override async def _process_response( self, diff --git a/src/kernel/lib/browser_routing/routing.py b/src/kernel/lib/browser_routing/routing.py index f005644f..0fbd5734 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/?$") @@ -198,6 +201,48 @@ 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. + """ + 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.append(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.append(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: """Whether a stale direct-to-VM auth failure can be retried on the control plane. diff --git a/tests/test_browser_routing.py b/tests/test_browser_routing.py index 756fbcda..ed4db6ac 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -3,7 +3,7 @@ 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 @@ -14,6 +14,7 @@ from kernel import ( Kernel, AsyncKernel, + APIConnectionError, AuthenticationError, InternalServerError, PermissionDeniedError, @@ -1259,3 +1260,96 @@ def test_indexed_multipart_body_flattens_only_given_values() -> None: "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 From 41ebedabdb271f775bb42f0ead4fc5daa24b4c63 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Wed, 2 Sep 2026 07:15:13 +0000 Subject: [PATCH 4/5] Run the route eviction hook before caller response hooks A response hook registered by the caller runs in registration order, so one that reads a failing 401/403 body or raises would skip the eviction hook appended after it and leave the stale route cached. Prepend the hook in both the sync and async installers; registration stays once per route cache. --- src/kernel/lib/browser_routing/routing.py | 7 +- tests/test_browser_routing.py | 110 ++++++++++++++++++++++ 2 files changed, 114 insertions(+), 3 deletions(-) diff --git a/src/kernel/lib/browser_routing/routing.py b/src/kernel/lib/browser_routing/routing.py index 0fbd5734..c9120e60 100644 --- a/src/kernel/lib/browser_routing/routing.py +++ b/src/kernel/lib/browser_routing/routing.py @@ -209,7 +209,8 @@ def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: Browse 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. + read, which keeps eviction independent of the body. It is prepended so that a + caller-supplied hook cannot pre-empt it by reading a failing body or raising. """ hooks = client.event_hooks.setdefault("response", []) if _has_eviction_hook(hooks, cache): @@ -220,7 +221,7 @@ def evict(response: httpx.Response) -> None: maybe_evict_browser_route_from_response(response, cache=cache) setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache) - hooks.append(evict) + hooks.insert(0, evict) def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient, *, cache: BrowserRouteCache) -> None: @@ -234,7 +235,7 @@ async def evict(response: httpx.Response) -> None: maybe_evict_browser_route_from_response(response, cache=cache) setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache) - hooks.append(evict) + hooks.insert(0, evict) def _has_eviction_hook(hooks: list[Any], cache: BrowserRouteCache) -> bool: diff --git a/tests/test_browser_routing.py b/tests/test_browser_routing.py index ed4db6ac..18cd9a20 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -1353,3 +1353,113 @@ def test_copied_client_registers_one_route_eviction_hook() -> None: 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 From fe23cc791631b414e68b83ebf9f0951ea379f9fc Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:05:52 +0000 Subject: [PATCH 5/5] Align routing tests with filesystem responses --- src/kernel/lib/browser_routing/routing.py | 6 ++++-- tests/test_browser_routing.py | 20 ++++++++++---------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/kernel/lib/browser_routing/routing.py b/src/kernel/lib/browser_routing/routing.py index c9120e60..df5f6567 100644 --- a/src/kernel/lib/browser_routing/routing.py +++ b/src/kernel/lib/browser_routing/routing.py @@ -209,8 +209,10 @@ def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: Browse 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. It is prepended so that a - caller-supplied hook cannot pre-empt it by reading a failing body or raising. + 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): diff --git a/tests/test_browser_routing.py b/tests/test_browser_routing.py index 18cd9a20..2ab6ad2a 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -751,7 +751,7 @@ def test_fs_read_file_routes_binary_response_from_vm(monkeypatch: pytest.MonkeyP 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(204) + return_value=httpx.Response(201) ) with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: _cache_browser(client) @@ -769,7 +769,7 @@ def test_fs_write_file_routes_binary_body_to_vm(monkeypatch: pytest.MonkeyPatch) @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(204)) + 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( @@ -905,7 +905,7 @@ def test_stale_direct_vm_jwt_replays_buffered_fs_body_on_control_plane( 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(204)) + 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") @@ -928,7 +928,7 @@ def test_stale_direct_vm_jwt_does_not_replay_streamed_fs_body( 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(204)) + 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): @@ -950,7 +950,7 @@ async def test_async_stale_direct_vm_jwt_does_not_replay_streamed_fs_body( 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(204)) + 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" @@ -978,7 +978,7 @@ def test_stale_direct_vm_jwt_replays_multipart_upload_on_control_plane( 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(204)) + 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") @@ -1094,7 +1094,7 @@ def test_stale_direct_vm_jwt_does_not_replay_multipart_that_cannot_rewind( 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(204)) + 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): @@ -1123,7 +1123,7 @@ async def test_async_stale_direct_vm_jwt_does_not_replay_multipart_that_cannot_r 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(204)) + 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 @@ -1152,7 +1152,7 @@ def test_stale_direct_vm_jwt_evicts_route_without_retries_for_buffered_body( 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(204)) + 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): @@ -1178,7 +1178,7 @@ async def test_async_stale_direct_vm_jwt_evicts_route_without_retries_for_stream 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(204)) + 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"