Pass use_alternative_endpoints to shell

Author: NickCao

packages/jumpstarter-cli-client/jumpstarter_cli_client/client_shell.py

@@ -22,6 +22,8 @@ def client_shell(config, selector: str, duration: timedelta, lease_name):
     with config.lease(selector=selector, lease_name=lease_name, duration=duration) as lease:
         with lease.serve_unix() as path:
             with lease.monitor():
-                exit_code = launch_shell(path, "remote", config.drivers.allow, config.drivers.unsafe)
+                exit_code = launch_shell(
+                    path, "remote", config.drivers.allow, config.drivers.unsafe, config.useAlternativeEndpoints
+                )
 
     sys.exit(exit_code)

packages/jumpstarter/jumpstarter/client/client.py

@@ -13,11 +13,18 @@
 
 
 @asynccontextmanager
-async def client_from_path(path: str, portal: BlockingPortal, stack: ExitStack, allow: list[str], unsafe: bool):
+async def client_from_path(
+    path: str,
+    portal: BlockingPortal,
+    stack: ExitStack,
+    allow: list[str],
+    unsafe: bool,
+    use_alternative_endpoints: bool = False,
+):
     async with grpc.aio.secure_channel(
         f"unix://{path}", grpc.local_channel_credentials(grpc.LocalConnectionType.UDS)
     ) as channel:
-        yield await client_from_channel(channel, portal, stack, allow, unsafe)
+        yield await client_from_channel(channel, portal, stack, allow, unsafe, use_alternative_endpoints)
 
 
 async def client_from_channel(

packages/jumpstarter/jumpstarter/common/utils.py

@@ -7,7 +7,7 @@
 
 from jumpstarter.client import client_from_path
 from jumpstarter.config.client import _allow_from_env
-from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JUMPSTARTER_HOST
+from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JMP_USE_ALTERNATIVE_ENDPOINTS, JUMPSTARTER_HOST
 from jumpstarter.driver import Driver
 from jumpstarter.exporter import Session
 
@@ -52,7 +52,16 @@ async def env_async(portal, stack):
 
     allow, unsafe = _allow_from_env()
 
-    async with client_from_path(host, portal, stack, allow=allow, unsafe=unsafe) as client:
+    use_alternative_endpoints = os.environ.get(JMP_USE_ALTERNATIVE_ENDPOINTS, "0") == "1"
+
+    async with client_from_path(
+        host,
+        portal,
+        stack,
+        allow=allow,
+        unsafe=unsafe,
+        use_alternative_endpoints=use_alternative_endpoints,
+    ) as client:
         try:
             yield client
         finally:
@@ -80,7 +89,13 @@ def env():
 PROMPT_CWD = "\\W"
 
 
-def launch_shell(host: str, context: str, allow: list[str], unsafe: bool) -> int:
+def launch_shell(
+    host: str,
+    context: str,
+    allow: list[str],
+    unsafe: bool,
+    use_alternative_endpoints: bool,
+) -> int:
     """Launch a shell with a custom prompt indicating the exporter type.
 
     Args:
@@ -98,6 +113,7 @@ def launch_shell(host: str, context: str, allow: list[str], unsafe: bool) -> int
         | {
             JUMPSTARTER_HOST: host,
             JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
+            JMP_USE_ALTERNATIVE_ENDPOINTS: "1" if use_alternative_endpoints else "0",
             "PS1": f"{ANSI_GRAY}{PROMPT_CWD} {ANSI_YELLOW}⚡{ANSI_WHITE}{context} {ANSI_YELLOW}➤{ANSI_RESET} ",
         },
     )

packages/jumpstarter/jumpstarter/config/env.py

@@ -7,3 +7,4 @@
 JMP_DRIVERS_ALLOW = "JMP_DRIVERS_ALLOW"
 JUMPSTARTER_HOST = "JUMPSTARTER_HOST"
 JMP_LEASE = "JMP_LEASE"
+JMP_USE_ALTERNATIVE_ENDPOINTS = "JMP_USE_ALTERNATIVE_ENDPOINTS"