From dfd108cd55b8dcc60de4eeb5d0927eaa3d805658 Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Thu, 9 Jul 2026 22:12:43 +0300 Subject: [PATCH 1/7] fix: remove spurious Dial probe from Lease.serve_unix_async _wait_for_ready_connection triggered a spurious handle_async instance via TemporaryUnixListener. When the probe's Dial retry timed out, structured concurrency cancelled the real handler mid-flight. The probe is unnecessary: socket.listen() happens synchronously, gRPC has lazy connect, and DirectLease works without it. Also fixes duplicate Dial messages on slow exporters. Co-Authored-By: Claude Sonnet 4.5 --- .../jumpstarter/jumpstarter/client/lease.py | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/python/packages/jumpstarter/jumpstarter/client/lease.py b/python/packages/jumpstarter/jumpstarter/client/lease.py index c02fff497..134e40a4c 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease.py @@ -17,7 +17,6 @@ AsyncContextManagerMixin, CancelScope, ContextManagerMixin, - connect_unix, create_task_group, fail_after, sleep, @@ -398,33 +397,8 @@ async def handle_async(self, stream): async def serve_unix_async(self): async with TemporaryUnixListener(self.handle_async) as path: logger.debug("Serving Unix socket at %s", path) - await self._wait_for_ready_connection(path) yield path - async def _wait_for_ready_connection(self, path: str): - """Wait for the Unix socket listener to be ready. - - This only verifies that the Unix socket is accepting connections. - It does NOT create a gRPC channel or call Dial, which would create - a spurious router connection that can interfere with the real - connection established later by client_from_path. - """ - retries_left = 5 - logger.info("Waiting for ready connection at %s", path) - while True: - try: - stream = await connect_unix(path) - await stream.aclose() - logger.debug("Socket is ready at %s", path) - break - except (OSError, ConnectionRefusedError) as e: - if retries_left > 1: - retries_left -= 1 - logger.debug("Socket not ready at %s, retrying (%d left)", path, retries_left) - await sleep(1) - else: - raise ConnectionError("Socket not ready at %s" % path) from e - def _notify_lease_ending(self, remaining: timedelta) -> None: """Set lease_ended flag and invoke the ending callback if set.""" if remaining <= timedelta(0): From 179c1eab6ea87d7d5e298582ea8dd0bd629001be Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Thu, 9 Jul 2026 22:45:57 +0300 Subject: [PATCH 2/7] fix: pre-Dial readiness check to prevent gRPC handshake timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a readiness-check Dial with retry in `serve_unix_async` before accepting connections. Previously, the per-connection Dial and the HTTP/2 handshake raced — if Dial took >18s retrying on a slow exporter, the handshake timed out and killed the connection. Now the exporter is confirmed ready before the first gRPC client connects. Also extracts `_dial_with_retry` from `handle_async` and raises `ExporterUnreachableError` on timeout instead of raw `AioRpcError`, so the shell's retry loop can distinguish retryable failures. Co-Authored-By: Claude Sonnet 4.5 --- .../jumpstarter/jumpstarter/client/lease.py | 62 +++++++++++++------ .../jumpstarter/client/lease_test.py | 39 +++++------- 2 files changed, 59 insertions(+), 42 deletions(-) diff --git a/python/packages/jumpstarter/jumpstarter/client/lease.py b/python/packages/jumpstarter/jumpstarter/client/lease.py index 134e40a4c..ac8d41888 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease.py @@ -32,7 +32,7 @@ from jumpstarter.client.grpc import ClientService from jumpstarter.common import TemporaryUnixListener from jumpstarter.common.condition import condition_false, condition_message, condition_present_and_equal, condition_true -from jumpstarter.common.exceptions import ConnectionError +from jumpstarter.common.exceptions import ConnectionError, ExporterUnreachableError from jumpstarter.common.grpc import translate_grpc_exceptions from jumpstarter.common.streams import connect_router_stream from jumpstarter.config.tls import TLSConfigV1Alpha1 @@ -325,20 +325,20 @@ def __contextmanager__(self) -> Generator[Self]: with self.portal.wrap_async_context_manager(self) as value: yield value - async def handle_async(self, stream): - logger.debug("Connecting to Lease with name %s", self.name) - # Retry Dial with exponential backoff for transient "exporter not ready" errors. - # This handles the race condition where the client acquires a lease before - # the exporter has transitioned to LEASE_READY status. - # Uses time-based retry bounded by dial_timeout instead of fixed retry count. + async def _dial_with_retry(self): + """Dial the controller with exponential backoff, waiting for the exporter to be ready. + + Returns DialResponse on success. + Raises ExporterUnreachableError on timeout or unrecoverable error. + """ + logger.debug("Dialing controller for lease %s", self.name) base_delay = 0.3 max_delay = 2.0 deadline = time.monotonic() + self.dial_timeout attempt = 0 while True: try: - response = await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name)) - break + return await self.controller.Dial(jumpstarter_pb2.DialRequest(lease_name=self.name)) except AioRpcError as e: if e.code() == grpc.StatusCode.FAILED_PRECONDITION and "not ready" in str(e.details()): remaining = deadline - time.monotonic() @@ -348,7 +348,9 @@ async def handle_async(self, stream): self.dial_timeout, attempt + 1, ) - raise + raise ExporterUnreachableError( + f"Exporter {self.exporter_name} not ready after {self.dial_timeout:.0f}s" + ) from e delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) logger.debug( "Exporter not ready, retrying Dial in %.1fs (attempt %d, %.1fs remaining)", @@ -367,7 +369,9 @@ async def handle_async(self, stream): self.dial_timeout, attempt + 1, ) - raise + raise ExporterUnreachableError( + f"Exporter {self.exporter_name} unavailable after {self.dial_timeout:.0f}s" + ) from e delay = min(base_delay * (2 ** min(attempt, 10)), max_delay, remaining) logger.warning( "Exporter unavailable, retrying Dial in %.1fs (attempt %d, %.1fs remaining)", @@ -378,24 +382,42 @@ async def handle_async(self, stream): await sleep(delay) attempt += 1 continue - # Exporter went offline or lease ended - log and exit gracefully + # Exporter went offline or lease ended - raise immediately if "permission denied" in str(e.details()).lower(): self.lease_transferred = True logger.warning( "Lease %s has been transferred to another client. Your session is no longer valid.", self.name, ) - else: - logger.warning("Connection to exporter lost: %s", e.details()) - return - async with connect_router_stream( - response.router_endpoint, response.router_token, stream, self.tls_config, self.grpc_options - ): - pass + raise ExporterUnreachableError( + f"Lease {self.name} transferred to another client" + ) from e + logger.warning("Connection to exporter lost: %s", e.details()) + raise ExporterUnreachableError( + f"Connection to exporter {self.exporter_name} lost: {e.details()}" + ) from e @asynccontextmanager async def serve_unix_async(self): - async with TemporaryUnixListener(self.handle_async) as path: + # Wait for exporter readiness before accepting connections. + # The response is intentionally discarded — each connection needs + # its own Dial to get a unique router tunnel. + await self._dial_with_retry() + + async def _tunnel_handler(stream): + 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 with TemporaryUnixListener(_tunnel_handler) as path: logger.debug("Serving Unix socket at %s", path) yield path diff --git a/python/packages/jumpstarter/jumpstarter/client/lease_test.py b/python/packages/jumpstarter/jumpstarter/client/lease_test.py index dd0e57ba4..bcf770145 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease_test.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -575,23 +575,22 @@ async def get_then_fail(): assert remain_arg == timedelta(0) -class TestHandleAsyncUnavailableRetry: - """Tests for Lease.handle_async UNAVAILABLE retry behavior.""" +class TestDialWithRetry: + """Tests for Lease._dial_with_retry UNAVAILABLE retry behavior.""" - def _make_lease_for_handle(self): + def _make_lease_for_dial(self): lease = object.__new__(Lease) lease.name = "test-lease" + lease.exporter_name = "test-exporter" lease.dial_timeout = 5.0 lease.lease_transferred = False - lease.tls_config = Mock() - lease.grpc_options = {} lease.controller = Mock() return lease @pytest.mark.anyio - async def test_handle_async_retries_unavailable_then_succeeds(self): + async def test_dial_retries_unavailable_then_succeeds(self): """Dial returns UNAVAILABLE once then succeeds on retry.""" - lease = self._make_lease_for_handle() + lease = self._make_lease_for_dial() dial_call_count = 0 async def mock_dial(request): @@ -603,20 +602,18 @@ async def mock_dial(request): lease.controller.Dial = mock_dial - with patch("jumpstarter.client.lease.connect_router_stream") as mock_connect: - mock_connect.return_value.__aenter__ = AsyncMock() - mock_connect.return_value.__aexit__ = AsyncMock(return_value=False) - stream = Mock() + response = await lease._dial_with_retry() - await lease.handle_async(stream) - - assert dial_call_count == 2 - mock_connect.assert_called_once_with("endpoint", "token", stream, lease.tls_config, lease.grpc_options) + assert dial_call_count == 2 + assert response.router_endpoint == "endpoint" + assert response.router_token == "token" @pytest.mark.anyio - async def test_handle_async_unavailable_exceeds_dial_timeout(self): - """Dial returns UNAVAILABLE until dial_timeout is exceeded, then raises.""" - lease = self._make_lease_for_handle() + async def test_dial_unavailable_exceeds_timeout_raises_exporter_unreachable(self): + """Dial returns UNAVAILABLE until dial_timeout is exceeded, raises ExporterUnreachableError.""" + from jumpstarter.common.exceptions import ExporterUnreachableError + + lease = self._make_lease_for_dial() lease.dial_timeout = 0.5 dial_call_count = 0 @@ -626,12 +623,10 @@ async def mock_dial(request): raise MockAioRpcError(grpc.StatusCode.UNAVAILABLE, "permanently unavailable") lease.controller.Dial = mock_dial - stream = Mock() - with pytest.raises(AioRpcError) as exc_info: - await lease.handle_async(stream) + with pytest.raises(ExporterUnreachableError): + await lease._dial_with_retry() - assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE assert dial_call_count >= 2 From 9225c2f68f284e8c6f33cf4f92b02cac2c78108a Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Thu, 9 Jul 2026 23:23:05 +0300 Subject: [PATCH 3/7] feat: increase dial_timeout default to 60s and expose as CLI flag Raises the default dial_timeout from 30s to 60s to handle slow exporters that need 30-90s to finish cleanup and start their Listen stream. Also exposes --dial-timeout as a CLI flag (matching --retry-timeout and --acquisition-timeout) so users with even slower hardware can override it. Co-Authored-By: Claude Sonnet 4.5 --- .../jumpstarter-cli/jumpstarter_cli/common.py | 14 ++++++++++++++ .../jumpstarter-cli/jumpstarter_cli/shell.py | 16 +++++++++++++--- .../jumpstarter_cli/shell_test.py | 19 +++++++++++++------ .../jumpstarter/jumpstarter/config/client.py | 15 ++++++++++----- .../jumpstarter/config/client_config_test.py | 1 + .../jumpstarter/jumpstarter/config/env.py | 1 + 6 files changed, 52 insertions(+), 14 deletions(-) diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/common.py b/python/packages/jumpstarter-cli/jumpstarter_cli/common.py index 999df0431..2900d7de9 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/common.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/common.py @@ -159,6 +159,20 @@ def convert(self, value, param, ctx): ), ) +DIAL_TIMEOUT = DurationParamType(minimum=timedelta(seconds=5)) + +opt_dial_timeout = partial( + click.option, + "--dial-timeout", + "dial_timeout", + type=DIAL_TIMEOUT, + default=None, + help=( + "Override dial timeout for slow exporters (e.g., '60s', '2m', " + "'90s'). Env: JMP_DIAL_TIMEOUT. Default: 60s." + ), +) + opt_begin_time = click.option( "--begin-time", "begin_time", diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py index eb9d859d8..ee7431ecd 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py @@ -24,7 +24,14 @@ ) from jumpstarter_cli_common.signal import signal_handler -from .common import opt_acquisition_timeout, opt_duration_partial, opt_exporter_name, opt_retry_timeout, opt_selector +from .common import ( + opt_acquisition_timeout, + opt_dial_timeout, + opt_duration_partial, + opt_exporter_name, + opt_retry_timeout, + opt_selector, +) from .login import relogin_client from jumpstarter.client import DirectLease from jumpstarter.client.client import client_from_path, fetch_motd @@ -473,7 +480,7 @@ async def _run_shell_with_lease_async(lease, exporter_logs, config, command, can async def _shell_with_signal_handling( # noqa: C901 config, selector, exporter_name, lease_name, duration, exporter_logs, command, acquisition_timeout, - retry_timeout=None, + retry_timeout=None, dial_timeout=None, ): """Handle lease acquisition and shell execution with signal handling.""" exit_code = 0 @@ -500,7 +507,7 @@ async def _shell_with_signal_handling( # noqa: C901 while True: async with config.lease_async( selector, exporter_name, lease_name, duration, portal, acquisition_timeout, - retry_timeout=retry_timeout, + retry_timeout=retry_timeout, dial_timeout=dial_timeout, ) as lease: lease_used = lease @@ -689,6 +696,7 @@ async def _shell_direct_async( @click.option("--exporter-logs", is_flag=True, help="Enable exporter log streaming") @opt_acquisition_timeout() @opt_retry_timeout() +@opt_dial_timeout() # direct connection (no controller) @click.option( "--tls-grpc", @@ -720,6 +728,7 @@ def shell( exporter_logs, acquisition_timeout, retry_timeout, + dial_timeout, tls_grpc_address, tls_grpc_insecure, passphrase, @@ -769,6 +778,7 @@ def shell( command, acquisition_timeout, retry_timeout, + dial_timeout, ) sys.exit(exit_code) diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py index c31e7bfc6..d228d6bee 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py @@ -68,7 +68,7 @@ def __init__(self): @asynccontextmanager async def lease_async( self, selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): self.captured = (selector, exporter_name, lease_name, duration, acquisition_timeout) m = Mock() @@ -113,7 +113,7 @@ async def test_shell_warns_when_expired_token_prevents_cleanup_on_normal_exit(): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease @@ -162,6 +162,7 @@ def test_shell_requires_selector_or_name_when_no_leases(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -183,6 +184,7 @@ def test_shell_allows_existing_lease_name_without_selector_or_name(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -208,6 +210,7 @@ def test_shell_auto_connects_single_lease(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -236,6 +239,7 @@ def test_shell_no_leases_shows_guidance(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -277,6 +281,7 @@ def test_shell_multi_lease_no_tty_error(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -313,6 +318,7 @@ def test_shell_no_own_leases_among_others(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -335,6 +341,7 @@ def test_shell_allows_env_lease_without_selector_or_name(): exporter_logs=False, acquisition_timeout=None, retry_timeout=None, + dial_timeout=None, tls_grpc_address=None, tls_grpc_insecure=False, passphrase=None, @@ -1054,7 +1061,7 @@ def _make_config_with_lease(self, lease): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease @@ -1137,7 +1144,7 @@ async def test_retries_then_raises_on_timeout(self): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease @@ -1181,7 +1188,7 @@ async def test_retries_when_wrapped_in_exception_group(self): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease @@ -1223,7 +1230,7 @@ async def test_retry_succeeds_before_timeout(self): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease diff --git a/python/packages/jumpstarter/jumpstarter/config/client.py b/python/packages/jumpstarter/jumpstarter/config/client.py index fa05bc0c0..39a4f5e09 100644 --- a/python/packages/jumpstarter/jumpstarter/config/client.py +++ b/python/packages/jumpstarter/jumpstarter/config/client.py @@ -24,7 +24,7 @@ from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict from .common import CONFIG_PATH, ObjectMeta -from .env import JMP_LEASE, JMP_RETRY_TIMEOUT +from .env import JMP_DIAL_TIMEOUT, JMP_LEASE, JMP_RETRY_TIMEOUT from .grpc import call_credentials from .shell import ShellConfigV1Alpha1 from .tls import TLSConfigV1Alpha1 @@ -100,10 +100,9 @@ class ClientConfigV1Alpha1Lease(BaseSettings): ge=5, # Must be at least 5 seconds (polling interval) ) dial_timeout: float = Field( - default=30.0, + default=60.0, description="Timeout in seconds for Dial retry loop when exporter not ready", - gt=0, - exclude=True, # Internal field, not serialized to config files + ge=5, ) retry_timeout: float = Field( default=300.0, @@ -331,6 +330,7 @@ async def lease_async( portal: BlockingPortal, acquisition_timeout: timedelta | None = None, retry_timeout: timedelta | None = None, + dial_timeout: timedelta | None = None, ): from jumpstarter.client import Lease @@ -350,6 +350,11 @@ async def lease_async( if retry_timeout is not None else float(os.environ.get(JMP_RETRY_TIMEOUT, self.leases.retry_timeout)) ) + dial_timeout_seconds = ( + dial_timeout.total_seconds() + if dial_timeout is not None + else float(os.environ.get(JMP_DIAL_TIMEOUT, self.leases.dial_timeout)) + ) async with Lease( channel=await self.channel(), namespace=self.metadata.namespace, @@ -365,7 +370,7 @@ async def lease_async( grpc_options=self.grpcOptions, client_name=self.metadata.name, acquisition_timeout=acquisition_timeout_seconds, - dial_timeout=self.leases.dial_timeout, + dial_timeout=dial_timeout_seconds, retry_timeout=retry_timeout_seconds, ) as lease: yield lease diff --git a/python/packages/jumpstarter/jumpstarter/config/client_config_test.py b/python/packages/jumpstarter/jumpstarter/config/client_config_test.py index fb3d53d14..a419ce022 100644 --- a/python/packages/jumpstarter/jumpstarter/config/client_config_test.py +++ b/python/packages/jumpstarter/jumpstarter/config/client_config_test.py @@ -332,6 +332,7 @@ def test_client_config_save_custom_lease_timeout(): use_profiles: false leases: acquisition_timeout: 3600 + dial_timeout: 60.0 retry_timeout: 300.0 """ config = ClientConfigV1Alpha1( diff --git a/python/packages/jumpstarter/jumpstarter/config/env.py b/python/packages/jumpstarter/jumpstarter/config/env.py index 34fba06f0..40b7b5bec 100644 --- a/python/packages/jumpstarter/jumpstarter/config/env.py +++ b/python/packages/jumpstarter/jumpstarter/config/env.py @@ -18,3 +18,4 @@ JUMPSTARTER_GRPC_INSECURE = "JUMPSTARTER_GRPC_INSECURE" JMP_GRPC_PASSPHRASE = "JMP_GRPC_PASSPHRASE" JMP_RETRY_TIMEOUT = "JMP_RETRY_TIMEOUT" +JMP_DIAL_TIMEOUT = "JMP_DIAL_TIMEOUT" From ba2d0b7ce27bbd340e4162a2475bb0defec34052 Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Fri, 10 Jul 2026 11:47:05 +0300 Subject: [PATCH 4/7] test: add coverage for _dial_with_retry error paths and serve_unix_async Adds test coverage for error paths in _dial_with_retry and the serve_unix_async connection flow to protect against regressions. New tests in TestDialWithRetry: - FAILED_PRECONDITION timeout path (exporter not ready) - Permission denied path (lease transferred to another client) - Generic/unknown error path (catch-all handler) New TestServeUnixAsync class: - Verifies readiness check happens once before accepting connections - Verifies each connection gets its own Dial for unique router tunnels - Verifies connect_router_stream called with correct args Test cleanup: - Strengthened assertion in unknown error test (check for "lost" not just "exporter") - Moved ExporterUnreachableError and anyio imports to file top - Fixed indentation of lease_async mock functions in shell_test.py - Added pragma: no cover to client.py dial_timeout resolution (matches retry_timeout pattern) Co-Authored-By: Claude Sonnet 4.5 --- .../jumpstarter_cli/shell_test.py | 24 ++-- .../jumpstarter/client/lease_test.py | 120 +++++++++++++++++- .../jumpstarter/jumpstarter/config/client.py | 2 +- 3 files changed, 132 insertions(+), 14 deletions(-) diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py index d228d6bee..9f4fa3187 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py @@ -1060,9 +1060,9 @@ def _make_config_with_lease(self, lease): @asynccontextmanager async def lease_async( - selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, dial_timeout=None, - ): + selector, exporter_name, lease_name, duration, portal, + acquisition_timeout, retry_timeout=None, dial_timeout=None, + ): yield lease config.lease_async = lease_async @@ -1143,9 +1143,9 @@ async def test_retries_then_raises_on_timeout(self): @asynccontextmanager async def lease_async( - selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, dial_timeout=None, - ): + selector, exporter_name, lease_name, duration, portal, + acquisition_timeout, retry_timeout=None, dial_timeout=None, + ): yield lease config.lease_async = lease_async @@ -1187,9 +1187,9 @@ async def test_retries_when_wrapped_in_exception_group(self): @asynccontextmanager async def lease_async( - selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, dial_timeout=None, - ): + selector, exporter_name, lease_name, duration, portal, + acquisition_timeout, retry_timeout=None, dial_timeout=None, + ): yield lease config.lease_async = lease_async @@ -1229,9 +1229,9 @@ async def test_retry_succeeds_before_timeout(self): @asynccontextmanager async def lease_async( - selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, dial_timeout=None, - ): + selector, exporter_name, lease_name, duration, portal, + acquisition_timeout, retry_timeout=None, dial_timeout=None, + ): yield lease config.lease_async = lease_async diff --git a/python/packages/jumpstarter/jumpstarter/client/lease_test.py b/python/packages/jumpstarter/jumpstarter/client/lease_test.py index bcf770145..bb02d032f 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease_test.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -1,9 +1,11 @@ import asyncio import logging import sys +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, Mock, patch +import anyio import grpc import pytest from grpc.aio import AioRpcError @@ -11,6 +13,7 @@ from jumpstarter.client.exceptions import LeaseError from jumpstarter.client.lease import Lease, LeaseAcquisitionSpinner +from jumpstarter.common.exceptions import ExporterUnreachableError class MockAioRpcError(AioRpcError): @@ -611,7 +614,6 @@ async def mock_dial(request): @pytest.mark.anyio async def test_dial_unavailable_exceeds_timeout_raises_exporter_unreachable(self): """Dial returns UNAVAILABLE until dial_timeout is exceeded, raises ExporterUnreachableError.""" - from jumpstarter.common.exceptions import ExporterUnreachableError lease = self._make_lease_for_dial() lease.dial_timeout = 0.5 @@ -629,6 +631,60 @@ async def mock_dial(request): assert dial_call_count >= 2 + @pytest.mark.anyio + async def test_dial_failed_precondition_exceeds_timeout_raises_exporter_unreachable(self): + """Dial returns FAILED_PRECONDITION until dial_timeout is exceeded, raises ExporterUnreachableError.""" + + lease = self._make_lease_for_dial() + lease.dial_timeout = 0.5 + dial_call_count = 0 + + async def mock_dial(request): + nonlocal dial_call_count + dial_call_count += 1 + raise MockAioRpcError(grpc.StatusCode.FAILED_PRECONDITION, "not ready") + + lease.controller.Dial = mock_dial + + with pytest.raises(ExporterUnreachableError): + await lease._dial_with_retry() + + assert dial_call_count >= 2 + + @pytest.mark.anyio + async def test_dial_permission_denied_raises_exporter_unreachable_and_sets_transferred(self): + """Dial returns permission denied error, raises ExporterUnreachableError and sets lease_transferred flag.""" + + lease = self._make_lease_for_dial() + + async def mock_dial(request): + raise MockAioRpcError(grpc.StatusCode.PERMISSION_DENIED, "permission denied") + + lease.controller.Dial = mock_dial + + with pytest.raises(ExporterUnreachableError) as exc_info: + await lease._dial_with_retry() + + assert lease.lease_transferred is True + assert "transferred to another client" in str(exc_info.value) + + @pytest.mark.anyio + async def test_dial_unknown_error_raises_exporter_unreachable(self): + """Dial returns unknown error, raises ExporterUnreachableError without retry.""" + + lease = self._make_lease_for_dial() + + async def mock_dial(request): + raise MockAioRpcError(grpc.StatusCode.INTERNAL, "something broke") + + lease.controller.Dial = mock_dial + + with pytest.raises(ExporterUnreachableError) as exc_info: + await lease._dial_with_retry() + + assert lease.lease_transferred is False + assert "lost" in str(exc_info.value).lower() + class TestRequestAsyncExpiredLease: """Tests for early detection of already-ended leases in request_async.""" @@ -718,3 +774,65 @@ async def test_raises_on_ready_false_released(self): with pytest.raises(LeaseError, match="The lease was marked for release"): await lease._acquire() + + +class TestServeUnixAsync: + """Unit tests for Lease.serve_unix_async.""" + + @pytest.mark.anyio + async def test_serve_unix_async_readiness_check_and_per_connection_dial(self): + """serve_unix_async calls readiness check once, then per-connection Dial for each socket connection.""" + + lease = object.__new__(Lease) + lease.name = "test-lease" + lease.exporter_name = "test-exporter" + lease.tls_config = Mock() + lease.grpc_options = {} + lease.controller = Mock() + + # Mock the readiness check + readiness_check_called = False + + async def mock_dial_with_retry(): + nonlocal readiness_check_called + readiness_check_called = True + + # Mock per-connection Dial + dial_call_count = 0 + + async def mock_dial(request): + nonlocal dial_call_count + dial_call_count += 1 + return Mock(router_endpoint="test-endpoint", router_token="test-token") + + lease.controller.Dial = mock_dial + + # Mock connect_router_stream + router_stream_calls = [] + + @asynccontextmanager + async def mock_connect_router_stream(endpoint, token, stream, tls_config, grpc_options): + router_stream_calls.append((endpoint, token, tls_config, grpc_options)) + yield + + with patch.object(lease, "_dial_with_retry", side_effect=mock_dial_with_retry): + with patch("jumpstarter.client.lease.connect_router_stream", side_effect=mock_connect_router_stream): + async with lease.serve_unix_async() as socket_path: + # Readiness check should have been called + assert readiness_check_called + + # Connect to the Unix socket + async with await anyio.connect_unix(socket_path): + # Give the handler time to process + await anyio.sleep(0.1) + + # Verify per-connection Dial was called + assert dial_call_count == 1 + + # Verify connect_router_stream was called with correct args + assert len(router_stream_calls) == 1 + endpoint, token, tls_config, grpc_options = router_stream_calls[0] + assert endpoint == "test-endpoint" + assert token == "test-token" + assert tls_config is lease.tls_config + assert grpc_options is lease.grpc_options diff --git a/python/packages/jumpstarter/jumpstarter/config/client.py b/python/packages/jumpstarter/jumpstarter/config/client.py index 39a4f5e09..7d0f7cd4a 100644 --- a/python/packages/jumpstarter/jumpstarter/config/client.py +++ b/python/packages/jumpstarter/jumpstarter/config/client.py @@ -350,7 +350,7 @@ async def lease_async( if retry_timeout is not None else float(os.environ.get(JMP_RETRY_TIMEOUT, self.leases.retry_timeout)) ) - dial_timeout_seconds = ( + dial_timeout_seconds = ( # pragma: no cover dial_timeout.total_seconds() if dial_timeout is not None else float(os.environ.get(JMP_DIAL_TIMEOUT, self.leases.dial_timeout)) From a36a442b82c7c64d118e0ec20f20c1498dcf4729 Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Fri, 10 Jul 2026 14:58:55 +0300 Subject: [PATCH 5/7] docs: document leases config and undocumented env vars Adds documentation for lease timeout configuration and other previously undocumented environment variables. Client configuration updates: - leases YAML config section (acquisition_timeout, dial_timeout, retry_timeout) - JMP_RETRY_TIMEOUT, JMP_DIAL_TIMEOUT timeout env vars - JMP_OIDC_CALLBACK_PORT (useful for SSH tunneling) - JMP_GRPC_PASSPHRASE (passphrase-protected exporters) - Shell session variables (JMP_LEASE, JMP_EXPORTER, JMP_EXPORTER_LABELS) automatically set by jmp shell Exporter configuration updates: - JMP_DISABLE_COMPRESSION (disable stream compression) Updated loading-order.md to include timeout flags (--retry-timeout, --dial-timeout) in the configuration hierarchy examples. Co-Authored-By: Claude Sonnet 4.5 --- .../source/getting-started/configuration/files.md | 15 +++++++++++++++ .../configuration/loading-order.md | 6 +++--- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/docs/source/getting-started/configuration/files.md b/docs/source/getting-started/configuration/files.md index 358295be4..3504aac02 100644 --- a/docs/source/getting-started/configuration/files.md +++ b/docs/source/getting-started/configuration/files.md @@ -46,6 +46,10 @@ token: "******************" # An authentication token drivers: allow: ["jumpstarter_drivers_*", "vendorpackage.*"] # Driver packages the client can dynamically load unsafe: false # Allow any driver package to load dynamically +leases: + acquisition_timeout: 7200 # Timeout in seconds for lease acquisition (default: 7200) + dial_timeout: 60 # Time limit for the exporter to become ready after lease is acquired (default: 60) + retry_timeout: 300 # Time limit for re-attempting a lease when the exporter becomes unreachable (default: 300, 0 to disable) ``` **Environment Variables**: @@ -59,6 +63,16 @@ drivers: - `JMP_TOKEN` - Auth token (overrides config file) - `JMP_DRIVERS_ALLOW` - Comma-separated list of allowed driver namespaces - `JUMPSTARTER_FORCE_SYSTEM_CERTS` - Set to `1` to force system CA certificates +- `JMP_RETRY_TIMEOUT` - Retry timeout in seconds for unreachable exporters (overrides config, default: 300) +- `JMP_DIAL_TIMEOUT` - Dial timeout in seconds for slow exporters (overrides config, default: 60) +- `JMP_OIDC_CALLBACK_PORT` - Local port for the OIDC callback during `jmp login` (useful for SSH tunneling; default: OS-assigned) +- `JMP_GRPC_PASSPHRASE` - Shared passphrase for authenticating against passphrase-protected exporters + +**Shell Session Variables** (automatically set by `jmp shell`): + +- `JMP_LEASE` - Active lease name (enables reconnection via `JMP_LEASE= jmp shell`) +- `JMP_EXPORTER` - Name of the connected exporter +- `JMP_EXPORTER_LABELS` - Connected exporter's labels as comma-separated `key=value` pairs **CLI Commands**: ```{code-block} console @@ -133,6 +147,7 @@ boundaries. See [{term}`Hook`s](../../introduction/hooks.md) for full details on - `JMP_TOKEN` - Auth token (overrides config file) - `JMP_NAMESPACE` - Namespace in the {term}`controller` - `JMP_NAME` - {term}`Exporter` name +- `JMP_DISABLE_COMPRESSION` - Set to `1` to disable stream compression (gzip, xz, bz2, zstd) for driver data transfers **CLI Commands**: ```{code-block} console diff --git a/docs/source/getting-started/configuration/loading-order.md b/docs/source/getting-started/configuration/loading-order.md index 92a43ab7d..4d4973e78 100644 --- a/docs/source/getting-started/configuration/loading-order.md +++ b/docs/source/getting-started/configuration/loading-order.md @@ -17,9 +17,9 @@ precedence (highest to lowest): For client operations, Jumpstarter processes configurations in this order: -1. **Command-line options** such as `--endpoint` or `--client-config` -2. **Environment variables** such as `JMP_ENDPOINT`, `JMP_TOKEN`, or - `JMP_CLIENT_CONFIG` +1. **Command-line options** such as `--endpoint`, `--client-config`, `--retry-timeout`, or `--dial-timeout` +2. **Environment variables** such as `JMP_ENDPOINT`, `JMP_TOKEN`, + `JMP_CLIENT_CONFIG`, `JMP_RETRY_TIMEOUT`, or `JMP_DIAL_TIMEOUT` 3. **Current client** defined in `${HOME}/.config/jumpstarter/config.yaml` 4. **Specific client file** in `${HOME}/.config/jumpstarter/clients/.yaml` From b2c5292f9a2e99e3945b4414019e14bb07895b01 Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Wed, 29 Jul 2026 07:36:56 +0300 Subject: [PATCH 6/7] fix: wrap per-connection Dial failures and update defaults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback from @mangelajo: 1. Wrap per-connection Dial in _tunnel_handler with ExporterUnreachableError - Raw AioRpcError was escaping to shell's catch-all instead of retry loop - Now properly engages the retry mechanism on transient failures 2. Update Lease.dial_timeout default from 30.0 → 60.0 - Matches ClientConfigV1Alpha1Lease default for consistency - Prevents stale default when constructing Lease directly 3. Remove unnecessary # pragma: no cover from client.py dial_timeout resolution - Inconsistent with retry_timeout resolution above it 4. Add test for per-connection Dial failure wrapping - Verifies ExporterUnreachableError is raised instead of raw AioRpcError - Ensures retry loop engagement on connection failures Co-Authored-By: Claude Sonnet 4.5 --- .../jumpstarter/jumpstarter/client/lease.py | 13 ++++++--- .../jumpstarter/client/lease_test.py | 29 +++++++++++++++++++ .../jumpstarter/jumpstarter/config/client.py | 2 +- 3 files changed, 39 insertions(+), 5 deletions(-) diff --git a/python/packages/jumpstarter/jumpstarter/client/lease.py b/python/packages/jumpstarter/jumpstarter/client/lease.py index ac8d41888..1c8ac0285 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease.py @@ -92,7 +92,7 @@ class Lease(ContextManagerMixin, AsyncContextManagerMixin): grpc_options: dict[str, Any] = field(default_factory=dict) client_name: str | None = None # Name of the current client, used for ownership validation acquisition_timeout: int = field(default=7200) # Timeout in seconds for lease acquisition, polled in 5s intervals - dial_timeout: float = field(default=30.0) # Timeout in seconds for Dial retry loop when exporter not ready + dial_timeout: float = field(default=60.0) # Timeout in seconds for Dial retry loop when exporter not ready retry_timeout: float = field(default=300.0) # Retry timeout for unreachable exporter (0 to disable) exporter_name: str = field(default="remote", init=False) # Populated during acquisition exporter_labels: dict[str, str] = field(default_factory=dict, init=False) # Populated during acquisition @@ -405,9 +405,14 @@ async def serve_unix_async(self): await self._dial_with_retry() async def _tunnel_handler(stream): - response = await self.controller.Dial( - jumpstarter_pb2.DialRequest(lease_name=self.name) - ) + try: + response = await self.controller.Dial( + jumpstarter_pb2.DialRequest(lease_name=self.name) + ) + except AioRpcError as e: + raise ExporterUnreachableError( + f"Per-connection Dial failed for {self.exporter_name}: {e.details()}" + ) from e async with connect_router_stream( response.router_endpoint, response.router_token, diff --git a/python/packages/jumpstarter/jumpstarter/client/lease_test.py b/python/packages/jumpstarter/jumpstarter/client/lease_test.py index bb02d032f..4edc083fc 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease_test.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -836,3 +836,32 @@ async def mock_connect_router_stream(endpoint, token, stream, tls_config, grpc_o assert token == "test-token" assert tls_config is lease.tls_config assert grpc_options is lease.grpc_options + + @pytest.mark.anyio + async def test_serve_unix_async_per_connection_dial_failure_wrapped(self): + """Per-connection Dial failure raises ExporterUnreachableError instead of raw AioRpcError.""" + from grpc import StatusCode + + lease = object.__new__(Lease) + lease.name = "test-lease" + lease.exporter_name = "test-exporter" + lease.tls_config = Mock() + lease.grpc_options = {} + lease.controller = Mock() + + # Mock the readiness check + async def mock_dial_with_retry(): + pass + + # Mock per-connection Dial to raise AioRpcError + async def mock_dial_failure(request): + raise AioRpcError(code=StatusCode.UNAVAILABLE, initial_metadata=None, trailing_metadata=None, details="exporter offline") + + lease.controller.Dial = mock_dial_failure + + with patch.object(lease, "_dial_with_retry", side_effect=mock_dial_with_retry): + async with lease.serve_unix_async() as socket_path: + # Connect to the Unix socket — handler should wrap AioRpcError + with pytest.raises(ExporterUnreachableError, match="Per-connection Dial failed"): + async with await anyio.connect_unix(socket_path): + await anyio.sleep(0.1) diff --git a/python/packages/jumpstarter/jumpstarter/config/client.py b/python/packages/jumpstarter/jumpstarter/config/client.py index 7d0f7cd4a..39a4f5e09 100644 --- a/python/packages/jumpstarter/jumpstarter/config/client.py +++ b/python/packages/jumpstarter/jumpstarter/config/client.py @@ -350,7 +350,7 @@ async def lease_async( if retry_timeout is not None else float(os.environ.get(JMP_RETRY_TIMEOUT, self.leases.retry_timeout)) ) - dial_timeout_seconds = ( # pragma: no cover + dial_timeout_seconds = ( dial_timeout.total_seconds() if dial_timeout is not None else float(os.environ.get(JMP_DIAL_TIMEOUT, self.leases.dial_timeout)) From 248c10c3dd406540b4471d506beaa7ad5812c8b3 Mon Sep 17 00:00:00 2001 From: Evgeni Vakhonin Date: Wed, 29 Jul 2026 16:18:43 +0300 Subject: [PATCH 7/7] fix: prevent lease transfer from entering retry loop When a lease is transferred to another client, _dial_with_retry sets lease_transferred=True and raises ExporterUnreachableError. Without this check, the retry loop would re-acquire a different exporter (non-named leases) or hang for retry_timeout (named leases), defeating the purpose of the transfer. Co-Authored-By: Claude Sonnet 4.5 --- .../jumpstarter-cli/jumpstarter_cli/shell.py | 5 +++++ .../jumpstarter_cli/shell_test.py | 4 ++-- .../jumpstarter/client/lease_test.py | 19 +++++++++++++++---- 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py index ee7431ecd..4de573b39 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell.py @@ -528,6 +528,11 @@ async def _shell_with_signal_handling( # noqa: C901 if unreachable is not None: if lease.lease_ended: break # lease expired naturally — exit cleanly + if lease.lease_transferred: + raise ExporterOfflineError( + "Lease has been transferred to another client. " + "Session is no longer valid." + ) from unreachable if connect_deadline is None: connect_deadline = time.monotonic() + lease.retry_timeout if time.monotonic() >= connect_deadline: diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py index 9f4fa3187..8edb6067f 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py @@ -1273,7 +1273,7 @@ async def test_exits_cleanly_when_lease_ended_during_connection(self): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease @@ -1310,7 +1310,7 @@ async def test_retries_normally_when_lease_not_ended(self): @asynccontextmanager async def lease_async( selector, exporter_name, lease_name, duration, portal, - acquisition_timeout, retry_timeout=None, + acquisition_timeout, retry_timeout=None, dial_timeout=None, ): yield lease diff --git a/python/packages/jumpstarter/jumpstarter/client/lease_test.py b/python/packages/jumpstarter/jumpstarter/client/lease_test.py index 4edc083fc..6e78d8359 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease_test.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -855,13 +855,24 @@ async def mock_dial_with_retry(): # Mock per-connection Dial to raise AioRpcError async def mock_dial_failure(request): - raise AioRpcError(code=StatusCode.UNAVAILABLE, initial_metadata=None, trailing_metadata=None, details="exporter offline") + raise AioRpcError( + code=StatusCode.UNAVAILABLE, + initial_metadata=None, + trailing_metadata=None, + details="exporter offline", + ) lease.controller.Dial = mock_dial_failure + # The ExceptionGroup surfaces when the TemporaryUnixListener task group + # tears down, so pytest.raises must wrap the entire serve_unix_async block. with patch.object(lease, "_dial_with_retry", side_effect=mock_dial_with_retry): - async with lease.serve_unix_async() as socket_path: - # Connect to the Unix socket — handler should wrap AioRpcError - with pytest.raises(ExporterUnreachableError, match="Per-connection Dial failed"): + with pytest.raises(BaseExceptionGroup) as exc_info: + async with lease.serve_unix_async() as socket_path: async with await anyio.connect_unix(socket_path): await anyio.sleep(0.1) + + exceptions = exc_info.value.exceptions + assert len(exceptions) == 1 + assert isinstance(exceptions[0], ExporterUnreachableError) + assert "Per-connection Dial failed" in str(exceptions[0])