Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions src/kernel/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
133 changes: 131 additions & 2 deletions src/kernel/lib/browser_routing/routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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/?$")
Expand All @@ -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()

Expand Down Expand Up @@ -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:
Expand Down
41 changes: 41 additions & 0 deletions src/kernel/lib/multipart.py
Original file line number Diff line number Diff line change
@@ -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
19 changes: 15 additions & 4 deletions src/kernel/resources/browsers/fs/fs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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", "<array>", "file"]])
extracted_files = extract_files(cast(Mapping[str, object], body), paths=[["files", "<array>", "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", "<array>", "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
Expand Down Expand Up @@ -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", "<array>", "file"]])
extracted_files = extract_files(cast(Mapping[str, object], body), paths=[["files", "<array>", "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", "<array>", "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
Expand Down
Loading
Loading