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

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

@@ -60,6 +60,7 @@ class DirectLease(ContextManagerMixin, AsyncContextManagerMixin):
 
     name: str = field(default="direct", init=False)
     exporter_name: str = field(default="direct", init=False)
+    exporter_labels: dict[str, str] = field(default_factory=dict, init=False)
     release: bool = field(default=False, init=False)
     lease_ended: bool = field(default=False, init=False)
 
@@ -92,6 +93,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
@@ -189,6 +191,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"):
@@ -229,6 +240,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,22 @@
 from anyio.from_thread import BlockingPortal, start_blocking_portal
 
 from jumpstarter.client import client_from_path
-from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JMP_GRPC_INSECURE, JMP_GRPC_PASSPHRASE, JUMPSTARTER_HOST
+from jumpstarter.config.env import (
+    JMP_DRIVERS_ALLOW,
+    JMP_EXPORTER,
+    JMP_EXPORTER_LABELS,
+    JMP_GRPC_INSECURE,
+    JMP_GRPC_PASSPHRASE,
+    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 +92,43 @@ 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 _build_common_env(
+    host: str,
+    allow: list[str],
+    unsafe: bool,
+    *,
+    lease=None,
+    insecure: bool = False,
+    passphrase: str | None = None,
+) -> dict[str, str]:
+    """Build the base environment dict for shell/command processes."""
+    env = os.environ | {
+        JUMPSTARTER_HOST: host,
+        JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
+        "_JMP_SUPPRESS_DRIVER_WARNINGS": "1",  # Already warned during client initialization
+    }
+    if insecure:
+        env = env | {JMP_GRPC_INSECURE: "1"}
+    if passphrase:
+        env = env | {JMP_GRPC_PASSPHRASE: passphrase}
+    if lease is not None:
+        env.update(_lease_env_vars(lease))
+    return env
+
+
 def launch_shell(
     host: str,
     context: str,
@@ -114,15 +159,9 @@ def launch_shell(
     shell = os.environ.get("SHELL", "bash")
     shell_name = os.path.basename(shell)
 
-    common_env = os.environ | {
-        JUMPSTARTER_HOST: host,
-        JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
-        "_JMP_SUPPRESS_DRIVER_WARNINGS": "1",  # Already warned during client initialization
-    }
-    if insecure:
-        common_env = common_env | {JMP_GRPC_INSECURE: "1"}
-    if passphrase:
-        common_env = common_env | {JMP_GRPC_PASSPHRASE: passphrase}
+    common_env = _build_common_env(
+        host, allow, unsafe, lease=lease, insecure=insecure, passphrase=passphrase
+    )
 
     if command:
         return _run_process(list(command), common_env, lease)

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

@@ -1,6 +1,11 @@
 import shutil
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
 
-from .utils import launch_shell
+import pytest
+
+from .utils import _build_common_env, _lease_env_vars, launch_shell
+from jumpstarter.utils.env import ExporterMetadata
 
 
 def test_launch_shell(tmp_path, monkeypatch):
@@ -22,3 +27,174 @@ def test_launch_shell(tmp_path, monkeypatch):
         use_profiles=False
     )
     assert exit_code == 1
+
+
+def test_launch_shell_sets_lease_env(tmp_path, monkeypatch):
+    env_output = tmp_path / "env_output.txt"
+    script = tmp_path / "capture_env.sh"
+    script.write_text(
+        f"#!/bin/sh\n"
+        f'echo "JMP_EXPORTER=$JMP_EXPORTER" >> {env_output}\n'
+        f'echo "JMP_LEASE=$JMP_LEASE" >> {env_output}\n'
+        f'echo "JMP_EXPORTER_LABELS=$JMP_EXPORTER_LABELS" >> {env_output}\n'
+    )
+    script.chmod(0o755)
+    monkeypatch.setenv("SHELL", str(script))
+    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
+    output = env_output.read_text()
+    assert "JMP_EXPORTER=my-exporter" in output
+    assert "JMP_LEASE=lease-123" in output
+    assert "board=rpi4" in output
+    assert "location=lab-1" in output
+
+
+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"}
+
+
+def test_exporter_metadata_from_env_ignores_empty_key(monkeypatch):
+    monkeypatch.setenv("JMP_EXPORTER", "board")
+    monkeypatch.setenv("JMP_EXPORTER_LABELS", "=value,valid=ok")
+
+    meta = ExporterMetadata.from_env()
+    assert meta.labels == {"valid": "ok"}
+
+
+def test_build_common_env_minimal():
+    env = _build_common_env("host.sock", ["driver1"], unsafe=False)
+    assert env["JUMPSTARTER_HOST"] == "host.sock"
+    assert env["JMP_DRIVERS_ALLOW"] == "driver1"
+    assert env["_JMP_SUPPRESS_DRIVER_WARNINGS"] == "1"
+    assert "JMP_GRPC_INSECURE" not in env
+    assert "JMP_GRPC_PASSPHRASE" not in env
+
+
+def test_build_common_env_unsafe():
+    env = _build_common_env("host.sock", ["driver1"], unsafe=True)
+    assert env["JMP_DRIVERS_ALLOW"] == "UNSAFE"
+
+
+def test_build_common_env_insecure():
+    env = _build_common_env("host.sock", ["*"], unsafe=False, insecure=True)
+    assert env["JMP_GRPC_INSECURE"] == "1"
+
+
+def test_build_common_env_passphrase():
+    env = _build_common_env("host.sock", ["*"], unsafe=False, passphrase="secret")
+    assert env["JMP_GRPC_PASSPHRASE"] == "secret"
+
+
+def test_build_common_env_empty_passphrase():
+    env = _build_common_env("host.sock", ["*"], unsafe=False, passphrase="")
+    assert "JMP_GRPC_PASSPHRASE" not in env
+
+
+def test_build_common_env_with_lease():
+    lease = SimpleNamespace(
+        exporter_name="exp1",
+        name="lease-1",
+        exporter_labels={"k": "v"},
+    )
+    env = _build_common_env("host.sock", ["*"], unsafe=False, lease=lease)
+    assert env["JMP_EXPORTER"] == "exp1"
+    assert env["JMP_LEASE"] == "lease-1"
+    assert env["JMP_EXPORTER_LABELS"] == "k=v"
+
+
+def test_lease_env_vars_basic():
+    lease = SimpleNamespace(
+        exporter_name="exp",
+        name="lease-x",
+        exporter_labels={"a": "1", "b": "2"},
+    )
+    env = _lease_env_vars(lease)
+    assert env["JMP_EXPORTER"] == "exp"
+    assert env["JMP_LEASE"] == "lease-x"
+    assert env["JMP_EXPORTER_LABELS"] == "a=1,b=2"
+
+
+def test_lease_env_vars_no_name_no_labels():
+    lease = SimpleNamespace(
+        exporter_name="exp",
+        name=None,
+        exporter_labels={},
+    )
+    env = _lease_env_vars(lease)
+    assert env["JMP_EXPORTER"] == "exp"
+    assert "JMP_LEASE" not in env
+    assert "JMP_EXPORTER_LABELS" not in env
+
+
+@pytest.mark.anyio
+async def test_fetch_exporter_labels_success():
+    from jumpstarter.client.lease import Lease
+
+    lease = object.__new__(Lease)
+    lease.exporter_name = "test-exporter"
+    lease.exporter_labels = {}
+
+    mock_exporter = MagicMock()
+    mock_exporter.labels = {"board": "rpi4", "env": "test"}
+    lease.svc = MagicMock()
+    lease.svc.GetExporter = AsyncMock(return_value=mock_exporter)
+
+    await lease._fetch_exporter_labels()
+
+    lease.svc.GetExporter.assert_called_once_with(name="test-exporter")
+    assert lease.exporter_labels == {"board": "rpi4", "env": "test"}
+
+
+@pytest.mark.anyio
+async def test_fetch_exporter_labels_failure():
+    from jumpstarter.client.lease import Lease
+
+    lease = object.__new__(Lease)
+    lease.exporter_name = "test-exporter"
+    lease.exporter_labels = {"stale": "data"}
+
+    lease.svc = MagicMock()
+    lease.svc.GetExporter = AsyncMock(side_effect=Exception("connection refused"))
+
+    await lease._fetch_exporter_labels()
+
+    assert lease.exporter_labels == {}

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"