Add exporter context env vars to jmp shell and env_with_metadata() helper

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

@@ -58,6 +58,7 @@ class Lease(ContextManagerMixin, AsyncContextManagerMixin):
     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
     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
     lease_ending_callback: Callable[[Self, timedelta], None] | None = field(
         default=None, init=False
     )  # Called when lease is ending
@@ -143,6 +144,15 @@ async def request_async(self):
 
         return await self._acquire()
 
+    async def _fetch_exporter_labels(self):
+        """Fetch the exporter's labels after lease acquisition."""
+        try:
+            exporter = await self.svc.GetExporter(name=self.exporter_name)
+            self.exporter_labels = exporter.labels
+        except Exception as e:
+            self.exporter_labels = {}
+            logger.warning("Could not fetch labels for exporter %s: %s", self.exporter_name, e)
+
     def _update_spinner_status(self, spinner, result):
         """Update spinner with appropriate status message based on lease conditions."""
         if condition_true(result.conditions, "Pending"):
@@ -183,6 +193,7 @@ async def _acquire(self):
                             logger.debug("Lease %s acquired", self.name)
                             spinner.update_status(f"Lease {self.name} acquired successfully!", force=True)
                             self.exporter_name = result.exporter
+                            await self._fetch_exporter_labels()
                             break
 
                         # lease unsatisfiable

python/packages/jumpstarter/jumpstarter/common/utils.py

@@ -10,14 +10,14 @@
 from anyio.from_thread import BlockingPortal, start_blocking_portal
 
 from jumpstarter.client import client_from_path
-from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JUMPSTARTER_HOST
+from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JMP_EXPORTER, JMP_EXPORTER_LABELS, JMP_LEASE, JUMPSTARTER_HOST
 from jumpstarter.exporter import Session
-from jumpstarter.utils.env import env
+from jumpstarter.utils.env import ExporterMetadata, env, env_with_metadata
 
 if TYPE_CHECKING:
     from jumpstarter.driver import Driver
 
-__all__ = ["env"]
+__all__ = ["ExporterMetadata", "env", "env_with_metadata"]
 
 
 @asynccontextmanager
@@ -84,6 +84,19 @@ def _run_process(
     return process.wait()
 
 
+def _lease_env_vars(lease) -> dict[str, str]:
+    """Extract environment variables from a lease object."""
+    env_vars: dict[str, str] = {}
+    env_vars[JMP_EXPORTER] = lease.exporter_name
+    if lease.name:
+        env_vars[JMP_LEASE] = lease.name
+    if lease.exporter_labels:
+        env_vars[JMP_EXPORTER_LABELS] = ",".join(
+            f"{k}={v}" for k, v in sorted(lease.exporter_labels.items())
+        )
+    return env_vars
+
+
 def launch_shell(
     host: str,
     context: str,
@@ -118,6 +131,9 @@ def launch_shell(
         "_JMP_SUPPRESS_DRIVER_WARNINGS": "1",  # Already warned during client initialization
     }
 
+    if lease is not None:
+        common_env.update(_lease_env_vars(lease))
+
     if command:
         return _run_process(list(command), common_env, lease)
 

python/packages/jumpstarter/jumpstarter/common/utils_test.py

@@ -1,6 +1,8 @@
 import shutil
+from types import SimpleNamespace
 
 from .utils import launch_shell
+from jumpstarter.utils.env import ExporterMetadata
 
 
 def test_launch_shell(tmp_path, monkeypatch):
@@ -22,3 +24,52 @@ def test_launch_shell(tmp_path, monkeypatch):
         use_profiles=False
     )
     assert exit_code == 1
+
+
+def test_launch_shell_sets_lease_env(tmp_path, monkeypatch):
+    monkeypatch.setenv("SHELL", shutil.which("env"))
+    lease = SimpleNamespace(
+        exporter_name="my-exporter",
+        name="lease-123",
+        exporter_labels={"board": "rpi4", "location": "lab-1"},
+        lease_ending_callback=None,
+    )
+    exit_code = launch_shell(
+        host=str(tmp_path / "test.sock"),
+        context="my-exporter",
+        allow=["*"],
+        unsafe=False,
+        use_profiles=False,
+        lease=lease,
+    )
+    assert exit_code == 0
+
+
+def test_exporter_metadata_from_env(monkeypatch):
+    monkeypatch.setenv("JMP_EXPORTER", "my-board")
+    monkeypatch.setenv("JMP_LEASE", "lease-abc")
+    monkeypatch.setenv("JMP_EXPORTER_LABELS", "board=rpi4,location=lab-1,team=qa")
+
+    meta = ExporterMetadata.from_env()
+    assert meta.name == "my-board"
+    assert meta.lease == "lease-abc"
+    assert meta.labels == {"board": "rpi4", "location": "lab-1", "team": "qa"}
+
+
+def test_exporter_metadata_from_env_empty(monkeypatch):
+    monkeypatch.delenv("JMP_EXPORTER", raising=False)
+    monkeypatch.delenv("JMP_LEASE", raising=False)
+    monkeypatch.delenv("JMP_EXPORTER_LABELS", raising=False)
+
+    meta = ExporterMetadata.from_env()
+    assert meta.name == ""
+    assert meta.lease is None
+    assert meta.labels == {}
+
+
+def test_exporter_metadata_from_env_labels_with_equals_in_value(monkeypatch):
+    monkeypatch.setenv("JMP_EXPORTER", "board")
+    monkeypatch.setenv("JMP_EXPORTER_LABELS", "key=val=123,other=ok")
+
+    meta = ExporterMetadata.from_env()
+    assert meta.labels == {"key": "val=123", "other": "ok"}

python/packages/jumpstarter/jumpstarter/config/env.py

@@ -8,6 +8,8 @@
 JMP_DRIVERS_ALLOW = "JMP_DRIVERS_ALLOW"
 JUMPSTARTER_HOST = "JUMPSTARTER_HOST"
 JMP_LEASE = "JMP_LEASE"
+JMP_EXPORTER = "JMP_EXPORTER"
+JMP_EXPORTER_LABELS = "JMP_EXPORTER_LABELS"
 
 JMP_DISABLE_COMPRESSION = "JMP_DISABLE_COMPRESSION"
 JMP_OIDC_CALLBACK_PORT = "JMP_OIDC_CALLBACK_PORT"

python/packages/jumpstarter/jumpstarter/utils/env.py

@@ -1,12 +1,38 @@
 import os
 from contextlib import ExitStack, asynccontextmanager, contextmanager
+from dataclasses import dataclass, field
 
 from anyio.from_thread import start_blocking_portal
 
 from jumpstarter.client import client_from_path
 from jumpstarter.common.exceptions import EnvironmentVariableNotSetError
 from jumpstarter.config.client import ClientConfigV1Alpha1Drivers
-from jumpstarter.config.env import JUMPSTARTER_HOST
+from jumpstarter.config.env import JMP_EXPORTER, JMP_EXPORTER_LABELS, JMP_LEASE, JUMPSTARTER_HOST
+
+
+@dataclass(frozen=True)
+class ExporterMetadata:
+    """Metadata about the exporter connected to the current shell session."""
+
+    name: str
+    labels: dict[str, str] = field(default_factory=dict)
+    lease: str | None = None
+
+    @classmethod
+    def from_env(cls) -> "ExporterMetadata":
+        """Build metadata from JMP_EXPORTER, JMP_EXPORTER_LABELS, and JMP_LEASE env vars."""
+        name = os.environ.get(JMP_EXPORTER, "")
+        lease = os.environ.get(JMP_LEASE) or None
+
+        labels: dict[str, str] = {}
+        raw_labels = os.environ.get(JMP_EXPORTER_LABELS, "")
+        if raw_labels:
+            for pair in raw_labels.split(","):
+                if "=" in pair:
+                    k, v = pair.split("=", 1)
+                    labels[k] = v
+
+        return cls(name=name, labels=labels, lease=lease)
 
 
 @asynccontextmanager
@@ -38,6 +64,19 @@ async def env_async(portal, stack):
                 client.close()
 
 
+@asynccontextmanager
+async def env_with_metadata_async(portal, stack):
+    """Provide a client and exporter metadata for an existing Jumpstarter shell.
+
+    Async version of env_with_metadata()
+
+    Yields a (client, ExporterMetadata) tuple. The metadata is read from environment
+    variables set by ``jmp shell``: JMP_EXPORTER, JMP_EXPORTER_LABELS, and JMP_LEASE.
+    """
+    async with env_async(portal, stack) as client:
+        yield client, ExporterMetadata.from_env()
+
+
 @contextmanager
 def env():
     """Provide a client for an existing JUMPSTARTER_HOST environment variable.
@@ -49,3 +88,23 @@ def env():
         with ExitStack() as stack:
             with portal.wrap_async_context_manager(env_async(portal, stack)) as client:
                 yield client
+
+
+@contextmanager
+def env_with_metadata():
+    """Provide a client and exporter metadata for an existing Jumpstarter shell.
+
+    This is useful when you need both the client and information about the connected
+    exporter (name, labels, lease ID).
+
+    Example::
+
+        with env_with_metadata() as (client, metadata):
+            print(metadata.name)
+            print(metadata.labels)
+            print(metadata.lease)
+    """
+    with start_blocking_portal() as portal:
+        with ExitStack() as stack:
+            with portal.wrap_async_context_manager(env_with_metadata_async(portal, stack)) as result:
+                yield result