fix(tunnel): poll established service readiness

Author: Reflex

src/runloop_api_client/lib/error_contract.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+import time
+import email.utils
 from typing import Mapping, cast
 from dataclasses import dataclass
 
@@ -37,6 +39,9 @@ def parse_retry_after(headers: httpx.Headers, body: object = None) -> float | No
     seconds = _number(headers.get("retry-after"))
     if seconds is not None:
         return seconds
+    retry_date = email.utils.parsedate_tz(headers.get("retry-after"))
+    if retry_date is not None:
+        return max(float(email.utils.mktime_tz(retry_date) - time.time()), 0)
     if isinstance(body, Mapping):
         payload = cast(Mapping[str, object], body)
         details = payload.get("details")

src/runloop_api_client/lib/tunnel_readiness.py

@@ -1,79 +1,167 @@
-"""Bounded retry helpers for tunnel service readiness."""
+"""Bounded polling of an established tunnel's service endpoint."""
 
 from __future__ import annotations
 
+import json
 import time
-import inspect
-from typing import TypeVar, Callable, Awaitable
+from typing import Mapping, Callable, Awaitable, cast
 
-from .._exceptions import APIStatusError
+import httpx
 
-T = TypeVar("T")
+from .._exceptions import APIError, APIStatusError, APITimeoutError, APIConnectionError
+from .error_contract import is_safe_transport_retry
 
 
-def _timeout(error: APIStatusError, *, port: int, path: str, timeout_seconds: float, attempts: int) -> None:
+def tunnel_url(*, api_host: str, tunnel_key: str, port: int, path: str = "/") -> str:
+    """Construct the established tunnel URL using the SDK's domain convention."""
+    if not 1 <= port <= 65535:
+        raise ValueError("port must be between 1 and 65535")
+    if not path.startswith("/"):
+        raise ValueError("path must start with '/'")
+    if not tunnel_key or not all(
+        character.isascii() and (character.isalnum() or character in "-_") for character in tunnel_key
+    ):
+        raise ValueError("tunnel_key contains characters that are unsafe in a tunnel hostname")
+    base_domain = api_host[4:] if api_host.startswith("api.") else api_host
+    return f"https://{port}-{tunnel_key}.tunnel.{base_domain}{path}"
+
+
+def tunnel_auth_headers(*, auth_mode: str, auth_token: str | None, request: httpx.Request) -> Mapping[str, str]:
+    """Return tunnel authentication without leaking the Runloop API bearer token."""
+    if auth_mode != "authenticated":
+        return {}
+    if auth_token:
+        return {"X-Runloop-Tunnel-Authorization": f"Bearer {auth_token}"}
+    error = APIConnectionError(
+        message="Authenticated tunnel is missing its tunnel authorization token.",
+        request=request,
+    )
+    error.code = "tunnel_authentication_required"
+    error.phase = "tunnel_readiness"
+    error.retryable = False
+    raise error
+
+
+def _status_error(response: httpx.Response, attempts: int) -> APIStatusError:
+    body: object = response.text
+    try:
+        body = cast(object, json.loads(response.text))
+    except (TypeError, ValueError):
+        pass
+    payload = cast(Mapping[str, object], body) if isinstance(body, dict) else None
+    body_message = payload.get("message") if payload is not None else None
+    message = (
+        body_message
+        if isinstance(body_message, str)
+        else f"Tunnel readiness check failed with HTTP status {response.status_code}."
+    )
+    error = APIStatusError(message, response=response, body=cast(object, body), attempts=attempts)
+    if error.phase == "api":
+        error.phase = "tunnel_readiness"
+    return error
+
+
+def _connection_error(error: httpx.HTTPError, request: httpx.Request, attempts: int) -> APIConnectionError:
+    if isinstance(error, httpx.TimeoutException):
+        return APITimeoutError(request=request, cause=error, attempts=attempts)
+    return APIConnectionError(request=request, cause=error, attempts=attempts)
+
+
+def _raise_timeout(error: APIError, *, port: int, path: str, timeout_seconds: float, attempts: int) -> None:
     message = f"Tunnel service was not ready for port {port} path {path!r} within {timeout_seconds:g} seconds."
     error.message = message
     error.args = (message,)
     error.attempts = attempts
     raise error
 
 
+def _retry_delay(error: APIError) -> float:
+    return error.retry_after if error.retry_after is not None else 0.5
+
+
+def _is_transient_status(error: APIStatusError) -> bool:
+    if error.code == "tunnel_unavailable":
+        return False
+    return error.code == "tunnel_service_not_ready" and error.response.headers.get("x-should-retry") != "false"
+
+
 def wait_for_tunnel_service(
-    operation: Callable[[], T],
+    request: Callable[[float], httpx.Response],
     *,
     port: int,
     path: str = "/",
     timeout_seconds: float = 30.0,
     clock: Callable[[], float] = time.monotonic,
     sleep: Callable[[float], None] = time.sleep,
-) -> T:
-    """Retry only ``tunnel_service_not_ready`` until a bounded deadline."""
+) -> None:
+    """Poll an established tunnel URL until it returns a successful response."""
     if timeout_seconds <= 0:
         raise ValueError("timeout_seconds must be greater than zero")
     deadline = clock() + timeout_seconds
     attempts = 0
+    failure: APIError | None = None
     while True:
+        remaining = deadline - clock()
+        if remaining <= 0:
+            if failure is None:
+                raise ValueError("tunnel readiness deadline expired before the first request")
+            _raise_timeout(failure, port=port, path=path, timeout_seconds=timeout_seconds, attempts=attempts)
         attempts += 1
         try:
-            return operation()
-        except APIStatusError as error:
-            error.attempts = attempts
-            if error.code != "tunnel_service_not_ready":
-                raise
-            remaining = deadline - clock()
-            if remaining <= 0 or attempts >= 1000:
-                _timeout(error, port=port, path=path, timeout_seconds=timeout_seconds, attempts=attempts)
-            delay = error.retry_after if error.retry_after is not None else 0.5
-            sleep(min(max(delay, 0), remaining))
+            response = request(remaining)
+            if response.is_success:
+                return
+            failure = _status_error(response, attempts)
+            if not _is_transient_status(failure):
+                raise failure
+        except httpx.HTTPError as cause:
+            request_object = cause.request
+            failure = _connection_error(cause, request_object, attempts)
+            if not is_safe_transport_retry(cause):
+                raise failure from cause
+
+        remaining = deadline - clock()
+        if remaining <= 0 or attempts >= 1000:
+            _raise_timeout(failure, port=port, path=path, timeout_seconds=timeout_seconds, attempts=attempts)
+        sleep(min(max(_retry_delay(failure), 0), remaining))
 
 
 async def async_wait_for_tunnel_service(
-    operation: Callable[[], Awaitable[T]],
+    request: Callable[[float], Awaitable[httpx.Response]],
     *,
     port: int,
     path: str = "/",
     timeout_seconds: float = 30.0,
     clock: Callable[[], float] = time.monotonic,
     sleep: Callable[[float], Awaitable[None]],
-) -> T:
+) -> None:
     """Async counterpart to :func:`wait_for_tunnel_service`."""
     if timeout_seconds <= 0:
         raise ValueError("timeout_seconds must be greater than zero")
     deadline = clock() + timeout_seconds
     attempts = 0
+    failure: APIError | None = None
     while True:
+        remaining = deadline - clock()
+        if remaining <= 0:
+            if failure is None:
+                raise ValueError("tunnel readiness deadline expired before the first request")
+            _raise_timeout(failure, port=port, path=path, timeout_seconds=timeout_seconds, attempts=attempts)
         attempts += 1
         try:
-            return await operation()
-        except APIStatusError as error:
-            error.attempts = attempts
-            if error.code != "tunnel_service_not_ready":
-                raise
-            remaining = deadline - clock()
-            if remaining <= 0 or attempts >= 1000:
-                _timeout(error, port=port, path=path, timeout_seconds=timeout_seconds, attempts=attempts)
-            delay = error.retry_after if error.retry_after is not None else 0.5
-            result = sleep(min(max(delay, 0), remaining))
-            if inspect.isawaitable(result):
-                await result
+            response = await request(remaining)
+            if response.is_success:
+                return
+            failure = _status_error(response, attempts)
+            if not _is_transient_status(failure):
+                raise failure
+        except httpx.HTTPError as cause:
+            request_object = cause.request
+            failure = _connection_error(cause, request_object, attempts)
+            if not is_safe_transport_retry(cause):
+                raise failure from cause
+
+        remaining = deadline - clock()
+        if remaining <= 0 or attempts >= 1000:
+            _raise_timeout(failure, port=port, path=path, timeout_seconds=timeout_seconds, attempts=attempts)
+        await sleep(min(max(_retry_delay(failure), 0), remaining))

src/runloop_api_client/sdk/async_devbox.py

@@ -8,6 +8,8 @@
 from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, Awaitable, cast
 from typing_extensions import Unpack, override
 
+import httpx
+
 from ..types import (
     DevboxView,
     TunnelView,
@@ -38,7 +40,7 @@
 from ..lib.polling import PollingConfig
 from ..types.devboxes import ExecutionUpdateChunk
 from .async_execution import AsyncExecution, _AsyncStreamingGroup
-from ..lib.tunnel_readiness import async_wait_for_tunnel_service
+from ..lib.tunnel_readiness import tunnel_url, tunnel_auth_headers, async_wait_for_tunnel_service
 from .async_execution_result import AsyncExecutionResult
 from ..types.devbox_execute_async_params import DevboxNiceExecuteAsyncParams
 from ..types.devboxes.devbox_logs_list_view import DevboxLogsListView
@@ -838,20 +840,36 @@ async def wait_for_tunnel_ready(
         path: str = "/",
         *,
         timeout_seconds: float = 30.0,
+        http_client: httpx.AsyncClient | None = None,
         clock: Callable[[], float] = time.monotonic,
         sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
         **params: Unpack[SDKDevboxEnableTunnelParams],
     ) -> TunnelView:
-        """Enable a tunnel, waiting through transient service readiness failures."""
+        """Enable a tunnel and poll the requested service until it is ready."""
         client = self._devbox._client.with_options(max_retries=0)
-        return await async_wait_for_tunnel_service(
-            lambda: client.devboxes.enable_tunnel(self._devbox.id, **params),
-            port=port,
-            path=path,
-            timeout_seconds=timeout_seconds,
-            clock=clock,
-            sleep=sleep,
+        enable_params: dict[str, Any] = dict(params)
+        enable_params.setdefault("timeout", timeout_seconds)
+        tunnel = await client.devboxes.enable_tunnel(self._devbox.id, **enable_params)
+        url = tunnel_url(api_host=client.base_url.host, tunnel_key=tunnel.tunnel_key, port=port, path=path)
+        headers = tunnel_auth_headers(
+            auth_mode=tunnel.auth_mode,
+            auth_token=tunnel.auth_token,
+            request=httpx.Request("GET", url),
         )
+        probe_client = http_client or httpx.AsyncClient(follow_redirects=True)
+        try:
+            await async_wait_for_tunnel_service(
+                lambda remaining: probe_client.get(url, headers=headers, timeout=remaining, follow_redirects=True),
+                port=port,
+                path=path,
+                timeout_seconds=timeout_seconds,
+                clock=clock,
+                sleep=sleep,
+            )
+        finally:
+            if http_client is None:
+                await probe_client.aclose()
+        return tunnel
 
     async def remove_tunnel(
         self,

src/runloop_api_client/sdk/devbox.py

@@ -8,6 +8,8 @@
 from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
 from typing_extensions import Unpack, override
 
+import httpx
+
 from ..types import (
     DevboxView,
     TunnelView,
@@ -40,7 +42,7 @@
 from ..lib.polling import PollingConfig
 from ..types.devboxes import ExecutionUpdateChunk
 from .execution_result import ExecutionResult
-from ..lib.tunnel_readiness import wait_for_tunnel_service
+from ..lib.tunnel_readiness import tunnel_url, tunnel_auth_headers, wait_for_tunnel_service
 from ..types.devbox_execute_async_params import DevboxNiceExecuteAsyncParams
 from ..types.devboxes.devbox_logs_list_view import DevboxLogsListView
 from ..types.devbox_async_execution_detail_view import DevboxAsyncExecutionDetailView
@@ -841,24 +843,40 @@ def wait_for_tunnel_ready(
         path: str = "/",
         *,
         timeout_seconds: float = 30.0,
+        http_client: httpx.Client | None = None,
         clock: Callable[[], float] = time.monotonic,
         sleep: Callable[[float], None] = time.sleep,
         **params: Unpack[SDKDevboxEnableTunnelParams],
     ) -> TunnelView:
-        """Enable a tunnel, waiting through transient service readiness failures.
+        """Enable a tunnel and poll the requested service until it is ready.
 
-        The generated client's own retries are disabled so this bounded helper
-        owns the deadline and preserves the final normalized error.
+        The readiness probe uses the tunnel-specific authorization token when
+        required and never forwards the Runloop API bearer token.
         """
         client = self._devbox._client.with_options(max_retries=0)
-        return wait_for_tunnel_service(
-            lambda: client.devboxes.enable_tunnel(self._devbox.id, **params),
-            port=port,
-            path=path,
-            timeout_seconds=timeout_seconds,
-            clock=clock,
-            sleep=sleep,
+        enable_params: dict[str, Any] = dict(params)
+        enable_params.setdefault("timeout", timeout_seconds)
+        tunnel = client.devboxes.enable_tunnel(self._devbox.id, **enable_params)
+        url = tunnel_url(api_host=client.base_url.host, tunnel_key=tunnel.tunnel_key, port=port, path=path)
+        headers = tunnel_auth_headers(
+            auth_mode=tunnel.auth_mode,
+            auth_token=tunnel.auth_token,
+            request=httpx.Request("GET", url),
         )
+        probe_client = http_client or httpx.Client(follow_redirects=True)
+        try:
+            wait_for_tunnel_service(
+                lambda remaining: probe_client.get(url, headers=headers, timeout=remaining, follow_redirects=True),
+                port=port,
+                path=path,
+                timeout_seconds=timeout_seconds,
+                clock=clock,
+                sleep=sleep,
+            )
+        finally:
+            if http_client is None:
+                probe_client.close()
+        return tunnel
 
     def remove_tunnel(
         self,