Validate on-demand sandbox substrate

Require DTS_SUBSTRATE to be injected as Sandbox or AcaSessionPool when constructing an on-demand sandbox worker.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: YunchuWang

durabletask-azuremanaged/durabletask/azuremanaged/preview/on_demand_sandbox/worker.py

@@ -45,6 +45,7 @@ def __init__(self) -> None:
         resolved_secure_channel = _resolve_secure_channel(resolved_host_address)
         resolved_token_credential = _resolve_token_credential()
         resolved_max_concurrent_activities = _resolve_max_concurrent_activities()
+        resolved_substrate = _resolve_substrate()
         concurrency_options = ConcurrencyOptions(
             maximum_concurrent_activity_work_items=resolved_max_concurrent_activities)
 
@@ -64,7 +65,7 @@ def __init__(self) -> None:
         self._on_demand_sandbox_worker_profile_id = _resolve_worker_profile_id()
         self._on_demand_sandbox_activity_names: list[str] = []
         self._on_demand_sandbox_max_activities = resolved_max_concurrent_activities
-        self._on_demand_sandbox_substrate = os.getenv("DTS_SUBSTRATE")
+        self._on_demand_sandbox_substrate = resolved_substrate
         self._on_demand_sandbox_dts_sandbox_identifier = os.getenv("DTS_SANDBOX_ID")
         self._on_demand_sandbox_heartbeat_interval_seconds = 2.0
         self._on_demand_sandbox_registration_stop = threading.Event()
@@ -211,6 +212,21 @@ def _resolve_token_credential() -> TokenCredential | None:
     return ManagedIdentityCredential(client_id=client_id.strip())
 
 
+def _resolve_substrate() -> str:
+    substrate = os.getenv("DTS_SUBSTRATE")
+    if not substrate:
+        raise ValueError(
+            "On-demand sandbox worker requires DTS_SUBSTRATE to be injected in the "
+            "sandbox environment.")
+
+    normalized = substrate.strip()
+    if normalized.lower() not in ("sandbox", "acasessionpool"):
+        raise ValueError(
+            "On-demand sandbox worker requires DTS_SUBSTRATE to be Sandbox or AcaSessionPool.")
+
+    return normalized
+
+
 def _resolve_max_concurrent_activities() -> int:
     value = os.getenv("DTS_ON_DEMAND_SANDBOX_MAX_ACTIVITIES")
     max_concurrent_activities = (

tests/durabletask-azuremanaged/test_on_demand_sandbox_extension.py

@@ -309,6 +309,7 @@ def test_on_demand_sandbox_activities_client_does_not_expose_worker_registration
 def test_on_demand_sandbox_worker_does_not_own_legacy_wakeup_server(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "http://localhost:8080")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
 
     worker = OnDemandSandboxWorker()
 
@@ -358,6 +359,7 @@ def OtherActivity(_ctx, value):
 def test_on_demand_sandbox_worker_stop_keeps_handle_for_still_running_registration_thread(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "http://localhost:8080")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
 
     class StillRunningThread:
         def __init__(self):
@@ -383,6 +385,7 @@ def is_alive(self):
 def test_on_demand_sandbox_worker_uses_scheduler_channel_without_credential(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
 
     worker = OnDemandSandboxWorker()
 
@@ -393,6 +396,7 @@ def test_on_demand_sandbox_worker_uses_scheduler_channel_without_credential(monk
 def test_on_demand_sandbox_worker_ignores_legacy_max_activities(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
     monkeypatch.delenv("DTS_ON_DEMAND_SANDBOX_MAX_ACTIVITIES", raising=False)
     monkeypatch.setenv("DTS_" + "SERVER" + "LESS_MAX_ACTIVITIES", "7")
 
@@ -404,6 +408,7 @@ def test_on_demand_sandbox_worker_ignores_legacy_max_activities(monkeypatch) ->
 def test_on_demand_sandbox_worker_tracks_active_activity_count_with_hooks(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
 
     worker = OnDemandSandboxWorker()
 
@@ -420,6 +425,7 @@ def test_on_demand_sandbox_worker_tracks_active_activity_count_with_hooks(monkey
 def test_on_demand_sandbox_worker_uses_managed_identity_credential_when_injected(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
     monkeypatch.setenv("DTS_AUTHENTICATION", "ManagedIdentity")
     monkeypatch.setenv("DTS_UMI_CLIENT_ID", "worker-client-id")
     monkeypatch.setattr(sandbox_worker, "ManagedIdentityCredential", _FakeManagedIdentityCredential)
@@ -434,6 +440,7 @@ def test_on_demand_sandbox_worker_uses_managed_identity_credential_when_injected
 def test_on_demand_sandbox_worker_requires_managed_identity_client_id_when_auth_enabled(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
     monkeypatch.setenv("DTS_AUTHENTICATION", "ManagedIdentity")
     monkeypatch.delenv("DTS_UMI_CLIENT_ID", raising=False)
 
@@ -448,6 +455,7 @@ def test_on_demand_sandbox_worker_requires_managed_identity_client_id_when_auth_
 def test_on_demand_sandbox_worker_requires_registered_activities(monkeypatch) -> None:
     monkeypatch.setenv("DTS_ENDPOINT", "http://localhost:8080")
     monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "Sandbox")
 
     worker = OnDemandSandboxWorker()
 
@@ -459,6 +467,32 @@ def test_on_demand_sandbox_worker_requires_registered_activities(monkeypatch) ->
         raise AssertionError("Expected missing registered activity names to fail.")
 
 
+def test_on_demand_sandbox_worker_requires_injected_substrate(monkeypatch) -> None:
+    monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
+    monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.delenv("DTS_SUBSTRATE", raising=False)
+
+    try:
+        OnDemandSandboxWorker()
+    except ValueError as ex:
+        assert "DTS_SUBSTRATE" in str(ex)
+    else:
+        raise AssertionError("Expected missing DTS_SUBSTRATE to fail.")
+
+
+def test_on_demand_sandbox_worker_rejects_invalid_substrate(monkeypatch) -> None:
+    monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
+    monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+    monkeypatch.setenv("DTS_SUBSTRATE", "ContainerApp")
+
+    try:
+        OnDemandSandboxWorker()
+    except ValueError as ex:
+        assert "Sandbox or AcaSessionPool" in str(ex)
+    else:
+        raise AssertionError("Expected invalid DTS_SUBSTRATE to fail.")
+
+
 class _RecordingChannel:
     def __init__(self) -> None:
         self.methods: list[str] = []