refactor: unify Dial and router retry into a single loop

Extract _dial_and_connect() to perform Dial + router connection as a
single atomic operation. This eliminates the duplicated Dial call that
was in the separate router retry block, addressing the code review
feedback about entangled and repeated code.

The single retry loop in handle_async now retries the full
_dial_and_connect() on transient errors, which naturally handles both
Dial failures and router connection failures with the same backoff
logic and always gets a fresh router token on retry.

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

Author: Ambient Code Bot

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

@@ -320,9 +320,21 @@ def __contextmanager__(self) -> Generator[Self]:
         grpc.StatusCode.INTERNAL,
     })
 
-    async def handle_async(self, stream):  # noqa: C901
+    async def _dial_and_connect(self, stream):
+        """Dial the controller and connect to the router stream.
+
+        Performs a single Dial + router connection attempt. Raises on failure
+        so the caller can decide whether to 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
+        ):
+            pass
+
+    async def handle_async(self, stream):
         logger.debug("Connecting to Lease with name %s", self.name)
-        # Retry Dial and router connection with exponential backoff for transient
+        # 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).
@@ -335,8 +347,8 @@ async def handle_async(self, stream):  # noqa: C901
         attempt = 0
         while True:
             try:
-                response = await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name))
-                break
+                await self._dial_and_connect(stream)
+                return
             except AioRpcError as e:
                 remaining = deadline - time.monotonic()
                 if e.code() == grpc.StatusCode.FAILED_PRECONDITION and "not ready" in str(e.details()):
@@ -349,7 +361,7 @@ async def handle_async(self, stream):  # noqa: C901
                         raise
                     delay = min(base_delay * (2**attempt), max_delay, remaining)
                     logger.debug(
-                        "Exporter not ready, retrying Dial in %.1fs (attempt %d, %.1fs remaining)",
+                        "Exporter not ready, retrying in %.1fs (attempt %d, %.1fs remaining)",
                         delay,
                         attempt + 1,
                         remaining,
@@ -361,15 +373,15 @@ async def handle_async(self, stream):  # noqa: C901
                 if e.code() in self._TRANSIENT_GRPC_CODES:
                     if remaining <= 0:
                         logger.warning(
-                            "Dial failed with transient error after %d attempts (%.1fs elapsed): %s",
+                            "Connection failed with transient error after %d attempts (%.1fs elapsed): %s",
                             attempt + 1,
                             self.dial_timeout,
                             e.details(),
                         )
                         return
                     delay = min(base_delay * (2**attempt), max_delay, remaining)
                     logger.info(
-                        "Dial failed with %s, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
+                        "Connection failed with %s, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
                         e.code().name,
                         delay,
                         attempt + 1,
@@ -389,60 +401,22 @@ async def handle_async(self, stream):  # noqa: C901
                 else:
                     logger.warning("Connection to exporter lost: %s", e.details())
                 return
-
-        # Connect to the router with retry for transient failures.
-        # After a successful Dial, the router endpoint may still be temporarily
-        # unreachable (e.g. after a tunnel drop). Retry the connection to give
-        # the network time to recover.
-        remaining = deadline - time.monotonic()
-        router_attempt = 0
-        while True:
-            try:
-                async with connect_router_stream(
-                    response.router_endpoint, response.router_token, stream, self.tls_config, self.grpc_options
-                ):
-                    return
-            except AioRpcError as e:
-                remaining = deadline - time.monotonic()
-                if e.code() in self._TRANSIENT_GRPC_CODES and remaining > 0:
-                    delay = min(base_delay * (2**router_attempt), max_delay, remaining)
-                    logger.info(
-                        "Router connection failed with %s, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
-                        e.code().name,
-                        delay,
-                        router_attempt + 1,
-                        remaining,
-                        e.details(),
-                    )
-                    await sleep(delay)
-                    router_attempt += 1
-                    # Re-dial to get a fresh router token since the old one may
-                    # have expired during the retry window
-                    try:
-                        response = await self.controller.Dial(
-                            jumpstarter_pb2.DialRequest(lease_name=self.name)
-                        )
-                    except AioRpcError:
-                        logger.debug("Re-dial failed during router retry, will retry from Dial")
-                    continue
-                logger.warning("Router connection failed: %s (code=%s)", e.details(), e.code().name)
-                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**router_attempt), max_delay, remaining)
+                    delay = min(base_delay * (2**attempt), max_delay, remaining)
                     logger.info(
-                        "Router connection failed with OSError, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
+                        "Connection failed with OSError, retrying in %.1fs (attempt %d, %.1fs remaining): %s",
                         delay,
-                        router_attempt + 1,
+                        attempt + 1,
                         remaining,
                         e,
                     )
                     await sleep(delay)
-                    router_attempt += 1
+                    attempt += 1
                     continue
-                logger.warning("Router connection failed: %s", e)
+                logger.warning("Connection failed: %s", e)
                 return
 
     @asynccontextmanager

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

@@ -582,12 +582,12 @@ def _make_lease_for_handle():
     return lease
 
 
-class TestHandleAsyncTransientDialRetry:
-    """Tests for transient gRPC error retry in handle_async Dial phase."""
+class TestHandleAsyncTransientRetry:
+    """Tests for transient gRPC error retry in handle_async (unified Dial + router loop)."""
 
     @pytest.mark.anyio
-    async def test_dial_retries_on_unavailable_then_succeeds(self):
-        """Dial should retry on UNAVAILABLE and succeed on the next attempt."""
+    async def test_retries_on_dial_unavailable_then_succeeds(self):
+        """Should retry on UNAVAILABLE from Dial and succeed on the next attempt."""
         lease = _make_lease_for_handle()
 
         dial_response = Mock(router_endpoint="ep", router_token="tok")
@@ -615,8 +615,8 @@ async def fake_router(*args, **kwargs):
         assert call_count == 2
 
     @pytest.mark.anyio
-    async def test_dial_transient_error_returns_after_timeout(self):
-        """Dial should give up and return when dial_timeout is exceeded."""
+    async def test_transient_error_returns_after_timeout(self):
+        """Should give up and return when dial_timeout is exceeded."""
         lease = _make_lease_for_handle()
         lease.dial_timeout = 0.0  # already expired
 
@@ -636,8 +636,8 @@ async def test_dial_transient_error_returns_after_timeout(self):
         [grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.ABORTED, grpc.StatusCode.INTERNAL],
         ids=["RESOURCE_EXHAUSTED", "ABORTED", "INTERNAL"],
     )
-    async def test_dial_retries_multiple_transient_codes(self, status_code):
-        """Dial should retry on RESOURCE_EXHAUSTED, ABORTED, INTERNAL."""
+    async def test_retries_multiple_transient_codes(self, status_code):
+        """Should retry on RESOURCE_EXHAUSTED, ABORTED, INTERNAL."""
         lease = _make_lease_for_handle()
         dial_response = Mock(router_endpoint="ep", router_token="tok")
         call_count = 0
@@ -663,13 +663,9 @@ async def fake_router(*args, **kwargs):
 
         assert call_count == 2, f"Expected 2 calls for {status_code}, got {call_count}"
 
-
-class TestHandleAsyncRouterRetry:
-    """Tests for router connection retry in handle_async."""
-
     @pytest.mark.anyio
-    async def test_router_retries_on_transient_error_then_succeeds(self):
-        """Router connection should retry on transient error, re-dial, then succeed."""
+    async def test_router_transient_error_retries_full_dial_and_connect(self):
+        """Router transient error should retry the full Dial + connect cycle."""
         lease = _make_lease_for_handle()
         dial_response = Mock(router_endpoint="ep", router_token="tok")
         lease.controller.Dial = AsyncMock(return_value=dial_response)
@@ -689,30 +685,30 @@ async def fake_router(*args, **kwargs):
                 await lease.handle_async(Mock())
 
         assert connect_count == 2
-        # Dial called once for initial + once for re-dial
+        # Dial is called fresh each attempt (unified loop)
         assert lease.controller.Dial.call_count == 2
 
     @pytest.mark.anyio
-    async def test_router_non_transient_error_returns_immediately(self):
-        """Router connection should not retry on non-transient errors."""
+    async def test_non_transient_error_returns_immediately(self):
+        """Non-transient errors should not be retried."""
         lease = _make_lease_for_handle()
         dial_response = Mock(router_endpoint="ep", router_token="tok")
         lease.controller.Dial = AsyncMock(return_value=dial_response)
 
         @asynccontextmanager
         async def fail_router(*args, **kwargs):
-            raise _make_aio_rpc_error(grpc.StatusCode.PERMISSION_DENIED, "no access")
+            raise _make_aio_rpc_error(grpc.StatusCode.NOT_FOUND, "not found")
             yield  # pragma: no cover
 
         with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fail_router):
             await lease.handle_async(Mock())
 
-        # Only the initial Dial, no re-dial
+        # Only one Dial attempt, no retry
         assert lease.controller.Dial.call_count == 1
 
     @pytest.mark.anyio
-    async def test_router_transient_error_returns_after_timeout(self):
-        """Router should give up when dial_timeout is exceeded."""
+    async def test_transient_router_error_returns_after_timeout(self):
+        """Should give up when dial_timeout is exceeded during router retries."""
         lease = _make_lease_for_handle()
         lease.dial_timeout = 0.0  # already expired
         dial_response = Mock(router_endpoint="ep", router_token="tok")
@@ -730,8 +726,8 @@ async def fail_router(*args, **kwargs):
         assert lease.controller.Dial.call_count == 1
 
     @pytest.mark.anyio
-    async def test_router_redial_failure_is_swallowed(self):
-        """When re-dial fails during router retry, the error is logged and retry continues."""
+    async def test_dial_failure_on_retry_is_retried_again(self):
+        """When Dial fails with a transient error during retry, it should keep retrying."""
         lease = _make_lease_for_handle()
         dial_response = Mock(router_endpoint="ep", router_token="tok")
 
@@ -741,11 +737,10 @@ async def dial_side_effect(req):
             nonlocal dial_count
             dial_count += 1
             if dial_count == 1:
-                return dial_response
+                return dial_response  # first Dial succeeds, router will fail
             if dial_count == 2:
-                # Re-dial fails
                 raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "re-dial failed")
-            return dial_response
+            return dial_response  # third Dial succeeds
 
         lease.controller.Dial = AsyncMock(side_effect=dial_side_effect)
 
@@ -755,22 +750,23 @@ async def dial_side_effect(req):
         async def fake_router(*args, **kwargs):
             nonlocal connect_count
             connect_count += 1
-            if connect_count <= 2:
+            if connect_count == 1:
                 raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "router fail")
             yield
 
         with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
             with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fake_router):
                 await lease.handle_async(Mock())
 
-        # Should have retried: connect fails, re-dial fails, connect fails again,
-        # re-dial succeeds, third connect succeeds
-        assert connect_count == 3
+        # Attempt 1: Dial OK -> router fails (UNAVAILABLE)
+        # Attempt 2: Dial fails (UNAVAILABLE) -> retried
+        # Attempt 3: Dial OK -> router OK
         assert dial_count == 3
+        assert connect_count == 2
 
     @pytest.mark.anyio
-    async def test_router_oserror_retries_then_succeeds(self):
-        """Router connection should retry on OSError, then succeed."""
+    async def test_oserror_retries_then_succeeds(self):
+        """OSError from router should retry the full Dial + connect cycle."""
         lease = _make_lease_for_handle()
         dial_response = Mock(router_endpoint="ep", router_token="tok")
         lease.controller.Dial = AsyncMock(return_value=dial_response)
@@ -790,10 +786,12 @@ async def fake_router(*args, **kwargs):
                 await lease.handle_async(Mock())
 
         assert connect_count == 2
+        # Dial called fresh each attempt
+        assert lease.controller.Dial.call_count == 2
 
     @pytest.mark.anyio
-    async def test_router_oserror_returns_after_timeout(self):
-        """Router should give up on OSError when dial_timeout is exceeded."""
+    async def test_oserror_returns_after_timeout(self):
+        """Should give up on OSError when dial_timeout is exceeded."""
         lease = _make_lease_for_handle()
         lease.dial_timeout = 0.0  # already expired
         dial_response = Mock(router_endpoint="ep", router_token="tok")