fix: address review feedback on retry logic

- Bound channel_ready_timeout by remaining dial_timeout deadline
- Change FAILED_PRECONDITION timeout from raise to return for consistency
- Extract _retry_delay() to deduplicate backoff computation
- Add type annotations to _dial_and_connect and handle_async
- Trim restating docstrings and comments
- Add tests for FAILED_PRECONDITION retry/timeout, PERMISSION_DENIED
  flag, and _retry_delay helper

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Ambient Code Bot

python/packages/jumpstarter/jumpstarter/client/lease.py

@@ -22,6 +22,7 @@
     fail_after,
     sleep,
 )
+from anyio.abc import SocketStream
 from anyio.from_thread import BlockingPortal
 from grpc.aio import AioRpcError, Channel
 from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
@@ -312,7 +313,8 @@ def __contextmanager__(self) -> Generator[Self]:
         with self.portal.wrap_async_context_manager(self) as value:
             yield value
 
-    # gRPC status codes that indicate transient network failures worth retrying.
+    # DEADLINE_EXCEEDED and CANCELLED are excluded: they indicate client-side
+    # timeout or cancellation, not server/network transients worth retrying.
     _TRANSIENT_GRPC_CODES = frozenset({
         grpc.StatusCode.UNAVAILABLE,
         grpc.StatusCode.RESOURCE_EXHAUSTED,
@@ -326,46 +328,53 @@ def __contextmanager__(self) -> Generator[Self]:
     # known to occur during tunnel reconnection.
     _TRANSIENT_UNKNOWN_MESSAGES = ("watch channel closed",)
 
-    async def _dial_and_connect(self, stream):
-        """Dial the controller and connect to the router stream.
+    @staticmethod
+    def _retry_delay(attempt: int, remaining: float, base: float = 0.3, cap: float = 5.0) -> float:
+        """Compute exponential-backoff delay, capped by *cap* and *remaining* time."""
+        return min(base * (2**attempt), cap, remaining)
 
-        Performs a single Dial + router connection attempt. Raises on failure
-        so the caller can decide whether to retry.
-        """
+    async def _dial_and_connect(
+        self, stream: SocketStream, channel_ready_timeout: float = 10.0
+    ) -> None:
+        """Single attempt; raises on failure for caller-driven retry."""
         response = await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name))
         async with connect_router_stream(
-            response.router_endpoint, response.router_token, stream, self.tls_config, self.grpc_options
+            response.router_endpoint,
+            response.router_token,
+            stream,
+            self.tls_config,
+            self.grpc_options,
+            channel_ready_timeout=channel_ready_timeout,
         ):
             pass
 
-    async def handle_async(self, stream):
+    async def handle_async(self, stream: SocketStream) -> None:
         logger.debug("Connecting to Lease with name %s", self.name)
-        # Retry Dial + router connection with exponential backoff for transient
-        # errors. This handles:
-        # 1. The race condition where the client acquires a lease before the
-        #    exporter has transitioned to LEASE_READY status (FAILED_PRECONDITION).
-        # 2. Transient network failures where the tunnel to the router drops and
-        #    needs to be re-established (UNAVAILABLE, etc.).
-        # Uses time-based retry bounded by dial_timeout instead of fixed retry count.
-        base_delay = 0.3
-        max_delay = 5.0
+        # Retry Dial + router connection with exponential backoff.
+        # Handles FAILED_PRECONDITION (exporter not yet ready), transient
+        # network errors (tunnel drops), and OSError (unreachable endpoint).
+        # All error paths return instead of raising because handle_async runs
+        # inside TemporaryUnixListener.serve's task group -- an unhandled
+        # exception would crash the listener and terminate sibling connections.
         deadline = time.monotonic() + self.dial_timeout
         attempt = 0
         while True:
+            remaining = deadline - time.monotonic()
+            channel_ready_timeout = max(min(10.0, remaining), 0.5)
             try:
-                await self._dial_and_connect(stream)
+                await self._dial_and_connect(stream, channel_ready_timeout=channel_ready_timeout)
                 return
             except AioRpcError as e:
                 remaining = deadline - time.monotonic()
                 if e.code() == grpc.StatusCode.FAILED_PRECONDITION and "not ready" in str(e.details()):
                     if remaining <= 0:
-                        logger.debug(
+                        logger.warning(
                             "Exporter not ready and dial timeout (%.1fs) exceeded after %d attempts",
                             self.dial_timeout,
                             attempt + 1,
                         )
-                        raise
-                    delay = min(base_delay * (2**attempt), max_delay, remaining)
+                        return
+                    delay = self._retry_delay(attempt, remaining)
                     logger.debug(
                         "Exporter not ready, retrying in %.1fs (attempt %d, %.1fs remaining)",
                         delay,
@@ -375,9 +384,6 @@ async def handle_async(self, stream):
                     await sleep(delay)
                     attempt += 1
                     continue
-                # Retry on transient network errors (e.g. tunnel to router dropped).
-                # Also retry UNKNOWN when the message matches a known transient
-                # tunnel teardown (e.g. "watch channel closed").
                 is_transient = e.code() in self._TRANSIENT_GRPC_CODES or (
                     e.code() == grpc.StatusCode.UNKNOWN
                     and any(msg in str(e.details()).lower() for msg in self._TRANSIENT_UNKNOWN_MESSAGES)
@@ -390,11 +396,8 @@ async def handle_async(self, stream):
                             self.dial_timeout,
                             e.details(),
                         )
-                        # Return instead of raising: handle_async runs inside
-                        # TemporaryUnixListener.serve's task group, so an
-                        # unhandled exception would crash the listener.
                         return
-                    delay = min(base_delay * (2**attempt), max_delay, remaining)
+                    delay = self._retry_delay(attempt, remaining)
                     logger.info(
                         "Connection failed with %s, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
                         e.code().name,
@@ -406,7 +409,6 @@ async def handle_async(self, stream):
                     await sleep(delay)
                     attempt += 1
                     continue
-                # Exporter went offline or lease ended - log and exit gracefully
                 if "permission denied" in str(e.details()).lower():
                     self.lease_transferred = True
                     logger.warning(
@@ -417,10 +419,9 @@ async def handle_async(self, stream):
                     logger.warning("Connection to exporter lost: %s", e.details())
                 return
             except OSError as e:
-                # OSError can occur when the router endpoint is unreachable
                 remaining = deadline - time.monotonic()
                 if remaining > 0:
-                    delay = min(base_delay * (2**attempt), max_delay, remaining)
+                    delay = self._retry_delay(attempt, remaining)
                     logger.info(
                         "Connection failed with OSError, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
                         delay,
@@ -432,7 +433,6 @@ async def handle_async(self, stream):
                     attempt += 1
                     continue
                 logger.warning("Connection failed: %s", e)
-                # Return instead of raising: see transient-error comment above.
                 return
 
     @asynccontextmanager

python/packages/jumpstarter/jumpstarter/client/lease_test.py

@@ -909,6 +909,87 @@ async def fake_router(*args, **kwargs):
             )
 
 
+    @pytest.mark.anyio
+    async def test_failed_precondition_not_ready_retries_then_succeeds(self):
+        """FAILED_PRECONDITION 'not ready' should retry and succeed on next attempt."""
+        lease = _make_lease_for_handle()
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        call_count = 0
+
+        async def dial_side_effect(req):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise _make_aio_rpc_error(
+                    grpc.StatusCode.FAILED_PRECONDITION, "exporter not ready"
+                )
+            return dial_response
+
+        lease.controller.Dial = AsyncMock(side_effect=dial_side_effect)
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            with patch("jumpstarter.client.lease.connect_router_stream") as mock_router:
+
+                @asynccontextmanager
+                async def fake_router(*args, **kwargs):
+                    yield
+
+                mock_router.side_effect = fake_router
+                await lease.handle_async(Mock())
+
+        assert call_count == 2
+
+    @pytest.mark.anyio
+    async def test_failed_precondition_returns_after_timeout(self):
+        """FAILED_PRECONDITION should return (not raise) when dial_timeout is exceeded."""
+        lease = _make_lease_for_handle()
+        lease.dial_timeout = 0.0  # already expired
+
+        lease.controller.Dial = AsyncMock(
+            side_effect=_make_aio_rpc_error(
+                grpc.StatusCode.FAILED_PRECONDITION, "exporter not ready"
+            ),
+        )
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            # Should return without raising
+            await lease.handle_async(Mock())
+
+        lease.controller.Dial.assert_called_once()
+
+    @pytest.mark.anyio
+    async def test_permission_denied_sets_lease_transferred(self):
+        """PERMISSION_DENIED should set lease_transferred = True."""
+        lease = _make_lease_for_handle()
+        assert lease.lease_transferred is False
+
+        lease.controller.Dial = AsyncMock(
+            side_effect=_make_aio_rpc_error(
+                grpc.StatusCode.PERMISSION_DENIED, "permission denied"
+            ),
+        )
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            await lease.handle_async(Mock())
+
+        assert lease.lease_transferred is True
+
+
+class TestRetryDelay:
+    """Tests for the _retry_delay static method."""
+
+    def test_basic_exponential(self):
+        assert Lease._retry_delay(0, 60.0) == pytest.approx(0.3)
+        assert Lease._retry_delay(1, 60.0) == pytest.approx(0.6)
+        assert Lease._retry_delay(2, 60.0) == pytest.approx(1.2)
+
+    def test_capped_by_max(self):
+        assert Lease._retry_delay(10, 60.0) == pytest.approx(5.0)
+
+    def test_capped_by_remaining(self):
+        assert Lease._retry_delay(0, 0.1) == pytest.approx(0.1)
+
+
 class TestTransientGrpcCodes:
     """Tests for the _TRANSIENT_GRPC_CODES class attribute."""