Separate sandbox management client from worker transport

Keep worker registration on an internal gRPC client while exposing only declaration management APIs through the public on-demand sandbox client.

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

Author: YunchuWang

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

@@ -233,8 +233,8 @@ def build_on_demand_sandbox_worker_heartbeat(active_activities_count: int) -> pb
             active_activities_count=active_activities_count))
 
 
-class OnDemandSandboxActivitiesClient:
-    """Client for Durable Task Scheduler on-demand sandbox activity management operations."""
+class _OnDemandSandboxActivitiesGrpcClient:
+    """Internal gRPC client for on-demand sandbox activity RPCs."""
 
     def __init__(
             self, *,
@@ -266,25 +266,60 @@ def close(self) -> None:
         if self._owns_channel:
             self._channel.close()
 
+    def declare_on_demand_sandbox_activities(
+            self,
+            declaration: pb.OnDemandSandboxActivityDeclaration) -> pb.OnDemandSandboxActivityDeclarationResult:
+        return self._stub.DeclareOnDemandSandboxActivities(declaration)
+
+    def remove_on_demand_sandbox_activity_declaration(
+            self,
+            worker_profile_id: str) -> pb.RemoveOnDemandSandboxActivityDeclarationResult:
+        return self._stub.RemoveOnDemandSandboxActivityDeclaration(
+            pb.RemoveOnDemandSandboxActivityDeclarationRequest(worker_profile_id=worker_profile_id))
+
+    def connect_on_demand_sandbox_activity_worker(
+            self,
+            messages: Iterable[pb.OnDemandSandboxActivityWorkerMessage]
+    ) -> pb.OnDemandSandboxActivityWorkerSessionResult:
+        return self._stub.ConnectOnDemandSandboxActivityWorker(messages)
+
+
+class OnDemandSandboxActivitiesClient:
+    """Client for Durable Task Scheduler on-demand sandbox activity management operations."""
+
+    def __init__(
+            self, *,
+            host_address: str,
+            taskhub: str,
+            token_credential: Optional[TokenCredential],
+            channel: Optional[grpc.Channel] = None,
+            secure_channel: bool = True,
+            interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
+            channel_options: Optional[GrpcChannelOptions] = None):
+        self._grpc_client = _OnDemandSandboxActivitiesGrpcClient(
+            host_address=host_address,
+            taskhub=taskhub,
+            token_credential=token_credential,
+            channel=channel,
+            secure_channel=secure_channel,
+            interceptors=interceptors,
+            channel_options=channel_options)
+
+    def close(self) -> None:
+        self._grpc_client.close()
+
     def enable_on_demand_sandbox_activities(self) -> None:
         """Declare all configured on-demand sandbox worker profiles with Durable Task Scheduler."""
         declarations = build_profile_on_demand_sandbox_activity_declarations()
         if not declarations:
             raise ValueError("No configured on-demand sandbox activities were found.")
 
         for declaration in declarations:
-            self._stub.DeclareOnDemandSandboxActivities(declaration)
+            self._grpc_client.declare_on_demand_sandbox_activities(declaration)
 
     def remove_on_demand_sandbox_activity_declaration(self, worker_profile_id: str) -> None:
         worker_profile_id = _normalize_required(worker_profile_id, "Worker profile ID is required.")
-        self._stub.RemoveOnDemandSandboxActivityDeclaration(
-            pb.RemoveOnDemandSandboxActivityDeclarationRequest(worker_profile_id=worker_profile_id))
-
-    def connect_on_demand_sandbox_activity_worker(
-            self,
-            messages: Iterable[pb.OnDemandSandboxActivityWorkerMessage]
-    ) -> pb.OnDemandSandboxActivityWorkerSessionResult:
-        return self._stub.ConnectOnDemandSandboxActivityWorker(messages)
+        self._grpc_client.remove_on_demand_sandbox_activity_declaration(worker_profile_id)
 
 
 def _normalize_optional_strings(values: Iterable[str]) -> list[str]:

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

@@ -12,7 +12,7 @@
 from durabletask.azuremanaged.preview.on_demand_sandbox.client import (
     DEFAULT_MAX_CONCURRENT_ACTIVITIES,
     DEFAULT_WORKER_PROFILE_ID,
-    OnDemandSandboxActivitiesClient,
+    _OnDemandSandboxActivitiesGrpcClient,
     build_on_demand_sandbox_worker_heartbeat,
     build_on_demand_sandbox_worker_start,
     resolve_activity_names,
@@ -118,7 +118,7 @@ def _run_on_demand_sandbox_registration_loop(self) -> None:
         retry_delay = 1.0
         while not self._on_demand_sandbox_registration_stop.is_set():
             try:
-                client = OnDemandSandboxActivitiesClient(
+                client = _OnDemandSandboxActivitiesGrpcClient(
                     host_address=self._on_demand_sandbox_host_address,
                     taskhub=self._on_demand_sandbox_taskhub,
                     token_credential=self._on_demand_sandbox_token_credential,

tests/durabletask-azuremanaged/test_on_demand_sandbox_extension.py

@@ -9,6 +9,7 @@
 import durabletask.azuremanaged.preview.on_demand_sandbox.client as sandbox_client
 import durabletask.azuremanaged.preview.on_demand_sandbox.worker as sandbox_worker
 from durabletask.azuremanaged.preview.on_demand_sandbox import OnDemandSandboxWorker
+from durabletask.azuremanaged.preview.on_demand_sandbox import OnDemandSandboxActivitiesClient
 from durabletask.azuremanaged.preview.on_demand_sandbox import OnDemandSandboxWorkerProfile
 from durabletask.azuremanaged.preview.on_demand_sandbox import OnDemandSandboxWorkerProfileOptions
 from durabletask.azuremanaged.preview.on_demand_sandbox import on_demand_sandbox_worker_profile
@@ -280,6 +281,10 @@ def test_on_demand_sandbox_worker_constructor_does_not_expose_runtime_contract()
     assert "_execute_activity" not in OnDemandSandboxWorker.__dict__
 
 
+def test_on_demand_sandbox_activities_client_does_not_expose_worker_registration_rpc() -> None:
+    assert not hasattr(OnDemandSandboxActivitiesClient, "connect_on_demand_sandbox_activity_worker")
+
+
 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")