Use base worker activity registration for sandboxes

Author: YunchuWang

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

@@ -5,12 +5,11 @@
 import random
 import threading
 
-from typing import Any, Iterator, Optional
+from typing import Iterator, Optional
 
 from azure.core.credentials import TokenCredential
 from azure.identity import ManagedIdentityCredential
 
-from durabletask import task
 from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb
 from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity
 from durabletask.azuremanaged.preview.sandboxes.helpers import resolve_activities
@@ -26,7 +25,6 @@
 import durabletask.internal.orchestrator_service_pb2 as worker_pb
 import durabletask.internal.shared as shared
 from durabletask.worker import (
-    ActivityWorkItemFilter,
     ConcurrencyOptions,
     WorkItemFilters,
 )
@@ -73,18 +71,6 @@ def __init__(self) -> None:
         self._sandbox_active_activities = 0
         self._sandbox_active_activities_lock = threading.Lock()
 
-    def add_activity(
-            self,
-            fn: task.Activity[Any, Any],
-            **kwargs: Any) -> str:
-        version = kwargs.pop("version", None)
-        if kwargs:
-            unexpected = next(iter(kwargs))
-            raise TypeError(f"Unexpected keyword argument: {unexpected}")
-        activity_name = super().add_activity(fn)
-        self._sandbox_activities.append(SandboxActivity(activity_name, version))
-        return activity_name
-
     def start(self) -> None:
         self._configure_sandbox_activity_filters()
         super().start()
@@ -103,17 +89,20 @@ def _on_activity_execution_completed(self, req: worker_pb.ActivityRequest) -> No
             self._sandbox_active_activities = max(0, self._sandbox_active_activities - 1)
 
     def _configure_sandbox_activity_filters(self) -> None:
-        activities = resolve_activities(self._sandbox_activities)
+        registered_filters = WorkItemFilters._from_registry(self._registry)  # pyright: ignore[reportPrivateUsage]
+        activities = resolve_activities([
+            SandboxActivity(activity_filter.name, version)
+            for activity_filter in registered_filters.activities
+            for version in (activity_filter.versions or [None])
+        ])
         if not activities:
             raise RuntimeError(
                 "Sandbox worker requires at least one registered activity before it can register.")
 
         self._sandbox_activities = activities
         self.use_work_item_filters(WorkItemFilters(
             orchestrations=[],
-            activities=[ActivityWorkItemFilter(
-                name=activity.name,
-                versions=[] if activity.version is None else [activity.version]) for activity in activities],
+            activities=registered_filters.activities,
             entities=[]))
 
     def _start_sandbox_registration(self) -> None:

tests/durabletask-azuremanaged/test_sandboxes_extension.py

@@ -436,6 +436,7 @@ def test_generated_stub_uses_sandbox_rpc_paths() -> None:
 def test_sandbox_worker_constructor_does_not_expose_runtime_contract() -> None:
     assert list(inspect.signature(SandboxWorker).parameters) == []
     assert "_execute_activity" not in SandboxWorker.__dict__
+    assert "add_activity" not in SandboxWorker.__dict__
 
 
 def test_sandbox_activities_client_does_not_expose_worker_registration_rpc() -> None:
@@ -472,8 +473,8 @@ def OtherActivity(_ctx, value):
         return value
 
     worker = SandboxWorker()
-    worker.add_activity(EnvActivity, version=None)
-    worker.add_activity(OtherActivity, version=None)
+    worker.add_activity(EnvActivity)
+    worker.add_activity(OtherActivity)
     worker._configure_sandbox_activity_filters()
     start = next(worker._registration_messages())