Merge pull request #164 from kernel/hypeship/direct-vm-fs-logs

Route browser fs and logs endpoints directly to the VM

Author: tnsardesai

src/kernel/_base_client.py

@@ -772,6 +772,9 @@ def _calculate_retry_timeout(
         timeout = sleep_seconds * jitter
         return timeout if timeout >= 0 else 0
 
+    def _should_retry_on_connection_error(self, _request: httpx.Request) -> bool:
+        return True
+
     def _should_retry(self, response: httpx.Response) -> bool:
         # Note: this is not a standard header
         should_retry_header = response.headers.get("x-should-retry")
@@ -1012,7 +1015,7 @@ def request(
             except httpx.TimeoutException as err:
                 log.debug("Encountered httpx.TimeoutException", exc_info=True)
 
-                if remaining_retries > 0:
+                if remaining_retries > 0 and self._should_retry_on_connection_error(request):
                     self._sleep_for_retry(
                         retries_taken=retries_taken,
                         max_retries=max_retries,
@@ -1026,7 +1029,7 @@ def request(
             except Exception as err:
                 log.debug("Encountered Exception", exc_info=True)
 
-                if remaining_retries > 0:
+                if remaining_retries > 0 and self._should_retry_on_connection_error(request):
                     self._sleep_for_retry(
                         retries_taken=retries_taken,
                         max_retries=max_retries,
@@ -1596,7 +1599,7 @@ async def request(
             except httpx.TimeoutException as err:
                 log.debug("Encountered httpx.TimeoutException", exc_info=True)
 
-                if remaining_retries > 0:
+                if remaining_retries > 0 and self._should_retry_on_connection_error(request):
                     await self._sleep_for_retry(
                         retries_taken=retries_taken,
                         max_retries=max_retries,
@@ -1610,7 +1613,7 @@ async def request(
             except Exception as err:
                 log.debug("Encountered Exception", exc_info=True)
 
-                if remaining_retries > 0:
+                if remaining_retries > 0 and self._should_retry_on_connection_error(request):
                     await self._sleep_for_retry(
                         retries_taken=retries_taken,
                         max_retries=max_retries,

src/kernel/_client.py

@@ -38,11 +38,18 @@
 from .lib.browser_routing.routing import (
     BrowserRouteCache,
     BrowserRoutingConfig,
-    strip_direct_vm_auth,
+    prepare_direct_vm_request,
     rewrite_direct_vm_options,
     browser_routing_config_from_env,
+    install_direct_vm_auth_stripping,
+    is_stale_direct_vm_auth_response,
     should_retry_stale_direct_vm_auth,
+    install_stale_direct_vm_auth_eviction,
+    install_async_direct_vm_auth_stripping,
     maybe_evict_browser_route_from_response,
+    should_retry_direct_vm_connection_error,
+    install_async_stale_direct_vm_auth_eviction,
+    direct_vm_request_body_is_known_unreplayable,
     maybe_populate_browser_route_cache_from_response,
 )
 
@@ -205,6 +212,8 @@ def __init__(
         )
         self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
         self._browser_routing = browser_routing_config_from_env()
+        install_direct_vm_auth_stripping(self._client)
+        install_stale_direct_vm_auth_eviction(self._client)
 
     @cached_property
     def deployments(self) -> DeploymentsResource:
@@ -369,13 +378,21 @@ def _prepare_options(self, options: Any) -> Any:
 
     @override
     def _prepare_request(self, request: httpx.Request) -> None:
-        strip_direct_vm_auth(request, cache=self.browser_route_cache)
+        prepare_direct_vm_request(request, cache=self.browser_route_cache)
+
+    @override
+    def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
+        return should_retry_direct_vm_connection_error(request)
 
     @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 direct_vm_request_body_is_known_unreplayable(response.request):
+            return False
+        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
@@ -594,6 +611,8 @@ def __init__(
         )
         self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
         self._browser_routing = browser_routing_config_from_env()
+        install_async_direct_vm_auth_stripping(self._client)
+        install_async_stale_direct_vm_auth_eviction(self._client)
 
     @cached_property
     def deployments(self) -> AsyncDeploymentsResource:
@@ -758,13 +777,21 @@ async def _prepare_options(self, options: Any) -> Any:
 
     @override
     async def _prepare_request(self, request: httpx.Request) -> None:
-        strip_direct_vm_auth(request, cache=self.browser_route_cache)
+        prepare_direct_vm_request(request, cache=self.browser_route_cache)
+
+    @override
+    def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
+        return should_retry_direct_vm_connection_error(request)
 
     @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 direct_vm_request_body_is_known_unreplayable(response.request):
+            return False
+        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

src/kernel/lib/browser_routing/raw_http.py

@@ -7,7 +7,7 @@
 import httpx
 
 from .util import sanitize_curl_raw_params
-from .routing import BrowserRoute
+from .routing import BrowserRoute, mark_direct_vm_headers, direct_vm_request_extensions
 from ..._types import Body, Timeout, NotGiven, not_given
 from ..._models import FinalRequestOptions
 
@@ -33,7 +33,7 @@ def request_via_browser_route(
         method=method.upper(),
         url=route.base_url.rstrip("/") + "/curl/raw",
         params=query,
-        headers=headers or {},
+        headers=mark_direct_vm_headers(headers),
         content=_normalize_binary_content(content),
         json_data=json,
         timeout=_normalize_timeout(timeout),
@@ -70,6 +70,7 @@ def stream_via_browser_route(
         headers=request_headers,
         content=_normalize_binary_content(content),
         timeout=_normalize_timeout(effective_timeout),
+        extensions=direct_vm_request_extensions(cache=parent.browser_route_cache),
     ) as response:
         yield response
 
@@ -93,7 +94,7 @@ async def async_request_via_browser_route(
         method=method.upper(),
         url=route.base_url.rstrip("/") + "/curl/raw",
         params=query,
-        headers=headers or {},
+        headers=mark_direct_vm_headers(headers),
         content=_normalize_binary_content(content),
         json_data=json,
         timeout=_normalize_timeout(timeout),
@@ -130,6 +131,7 @@ async def async_stream_via_browser_route(
         headers=request_headers,
         content=_normalize_binary_content(content),
         timeout=_normalize_timeout(effective_timeout),
+        extensions=direct_vm_request_extensions(cache=parent.browser_route_cache),
     ) as response:
         yield response