Pass use_alternative_endpoints to shell

Author: NickCao

packages/jumpstarter-cli/jumpstarter_cli/shell.py

@@ -31,7 +31,13 @@ def shell(config, lease_name, selector, duration):
             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,
+                            use_alternative_endpoints=config.use_alternative_endpoints,
+                        )
 
             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:
@@ -103,6 +118,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"