Fix on-demand sandbox CI checks

Use package-internal public helper names for cross-module calls, add strict typing around generated gRPC transport methods, and keep typed dataclass default factories without extra helper functions.

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

Author: YunchuWang

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

@@ -6,9 +6,9 @@
 import grpc
 from azure.core.credentials import TokenCredential
 
-from durabletask.azuremanaged.preview.on_demand_sandbox.helpers import _normalize_required
+from durabletask.azuremanaged.preview.on_demand_sandbox.helpers import normalize_required
 from durabletask.azuremanaged.preview.on_demand_sandbox.declarations import (
-    _build_profile_on_demand_sandbox_activity_declarations,
+    build_profile_on_demand_sandbox_activity_declarations,
 )
 from durabletask.azuremanaged.preview.on_demand_sandbox.transport import (
     OnDemandSandboxActivitiesGrpcTransport,
@@ -43,13 +43,13 @@ def close(self) -> None:
 
     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()
+        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._transport.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.")
+        worker_profile_id = normalize_required(worker_profile_id, "Worker profile ID is required.")
         self._transport.remove_on_demand_sandbox_activity_declaration(worker_profile_id)

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

@@ -3,13 +3,13 @@
 
 from dataclasses import dataclass, field
 from decimal import Decimal, InvalidOperation
-from typing import Callable, Iterable, Optional
+from typing import Any, Callable, Iterable, Optional
 
 from durabletask import task
 from durabletask.azuremanaged.internal import on_demand_sandbox_activities_service_pb2 as pb
 from durabletask.azuremanaged.preview.on_demand_sandbox.helpers import (
-    _normalize_required,
-    _resolve_activity_names,
+    normalize_required,
+    resolve_activity_names,
 )
 
 
@@ -32,17 +32,17 @@ class OnDemandSandboxWorkerProfileOptions:
     scheduler_managed_identity_client_id: Optional[str] = None
     cpu: str = DEFAULT_CPU
     memory: str = DEFAULT_MEMORY
-    environment_variables: dict[str, str] = field(default_factory=dict)
+    environment_variables: dict[str, str] = field(default_factory=dict[str, str])
     max_concurrent_activities: int = DEFAULT_MAX_CONCURRENT_ACTIVITIES
-    entrypoint: list[str] = field(default_factory=list)
-    cmd: list[str] = field(default_factory=list)
-    activity_names: list[str] = field(default_factory=list)
+    entrypoint: list[str] = field(default_factory=list[str])
+    cmd: list[str] = field(default_factory=list[str])
+    activity_names: list[str] = field(default_factory=list[str])
 
-    def add_activity(self, activity: str | Callable) -> None:
+    def add_activity(self, activity: str | Callable[..., Any]) -> None:
         """Add an activity to the on-demand sandbox worker profile declaration."""
         activity_name = task.get_name(activity) if callable(activity) else activity
         self.activity_names.append(
-            _normalize_required(activity_name, "On-demand sandbox activity name is required."))
+            normalize_required(activity_name, "On-demand sandbox activity name is required."))
 
 
 class OnDemandSandboxWorkerProfile:
@@ -57,7 +57,7 @@ def configure(self, options: OnDemandSandboxWorkerProfileOptions) -> None:
 
 def on_demand_sandbox_worker_profile(worker_profile_id: str) -> Callable[[type], type]:
     """Declare an on-demand sandbox worker profile using a decorated marker class."""
-    normalized_profile = _normalize_required(worker_profile_id, "On-demand sandbox worker profile ID is required.")
+    normalized_profile = normalize_required(worker_profile_id, "On-demand sandbox worker profile ID is required.")
 
     def decorator(cls: type) -> type:
         if normalized_profile in _worker_profiles:
@@ -73,7 +73,7 @@ def decorator(cls: type) -> type:
         if callable(configure):
             configure(options)
 
-        if not _resolve_activity_names(options.activity_names):
+        if not resolve_activity_names(options.activity_names):
             raise ValueError(
                 f"On-demand sandbox worker profile '{normalized_profile}' must declare at least one activity.")
 
@@ -86,7 +86,7 @@ def decorator(cls: type) -> type:
 def _build_on_demand_sandbox_activity_declaration(
         *,
         activity_names: str | Iterable[str],
-        scheduler_managed_identity_client_id: str,
+        scheduler_managed_identity_client_id: Optional[str],
         worker_profile_id: str = DEFAULT_WORKER_PROFILE_ID,
         container_image: Optional[str] = None,
         image_pull_managed_identity_client_id: Optional[str] = None,
@@ -103,7 +103,7 @@ def _build_on_demand_sandbox_activity_declaration(
             such as "myregistry.azurecr.io/workers/hello:1.0" or
             "myregistry.azurecr.io/workers/hello@sha256:0123456789abcdef...".
     """
-    resolved_activity_names = _resolve_activity_names(activity_names)
+    resolved_activity_names = resolve_activity_names(activity_names)
     if not resolved_activity_names:
         raise ValueError("On-demand sandbox activity declaration requires at least one activity name.")
 
@@ -113,16 +113,16 @@ def _build_on_demand_sandbox_activity_declaration(
     if max_concurrent_activities <= 0:
         raise ValueError("On-demand sandbox activity max concurrent activities must be greater than zero.")
 
-    image_ref = _normalize_required(
+    image_ref = normalize_required(
         container_image,
         "On-demand sandbox activity image metadata requires a container image reference like "
         "'myregistry.azurecr.io/workers/hello:1.0' or "
         "'myregistry.azurecr.io/workers/hello@sha256:...'.")
 
-    resolved_scheduler_managed_identity_client_id = _normalize_required(
+    resolved_scheduler_managed_identity_client_id = normalize_required(
         scheduler_managed_identity_client_id,
         "On-demand sandbox activity declaration requires the managed identity client ID workers use to connect to Durable Task Scheduler.")
-    resolved_image_pull_managed_identity_client_id = _normalize_required(
+    resolved_image_pull_managed_identity_client_id = normalize_required(
         image_pull_managed_identity_client_id,
         "On-demand sandbox activity declaration requires the managed identity client ID ADC uses to pull the worker image.")
 
@@ -146,19 +146,19 @@ def _build_on_demand_sandbox_activity_declaration(
     return declaration
 
 
-def _build_profile_on_demand_sandbox_activity_declarations() -> list[pb.OnDemandSandboxActivityDeclaration]:
+def build_profile_on_demand_sandbox_activity_declarations() -> list[pb.OnDemandSandboxActivityDeclaration]:
     """Build on-demand sandbox declarations from worker profile configuration."""
     declarations: list[pb.OnDemandSandboxActivityDeclaration] = []
     activity_owners: dict[str, str] = {}
     for profile in _worker_profiles.values():
-        activity_names = _resolve_activity_names(profile.activity_names)
+        activity_names = resolve_activity_names(profile.activity_names)
 
         for activity_name in activity_names:
             existing_profile = activity_owners.get(activity_name)
             if existing_profile and existing_profile != profile.worker_profile_id:
                 raise ValueError(
                     f"On-demand sandbox activity '{activity_name}' is assigned to both worker profile "
-                f"'{existing_profile}' and '{profile.worker_profile_id}'.")
+                    f"'{existing_profile}' and '{profile.worker_profile_id}'.")
             activity_owners[activity_name] = profile.worker_profile_id
 
         declarations.append(_build_on_demand_sandbox_activity_declaration(
@@ -177,7 +177,7 @@ def _build_profile_on_demand_sandbox_activity_declarations() -> list[pb.OnDemand
     return declarations
 
 
-def _build_on_demand_sandbox_worker_start(
+def build_on_demand_sandbox_worker_start(
         *,
         taskhub: str,
         worker_profile_id: str,
@@ -194,7 +194,7 @@ def _build_on_demand_sandbox_worker_start(
     if max_activities_count <= 0:
         raise ValueError("On-demand sandbox activity worker max activity count must be greater than zero.")
 
-    resolved_activity_names = _resolve_activity_names(activity_names)
+    resolved_activity_names = resolve_activity_names(activity_names)
     if not resolved_activity_names:
         raise ValueError("On-demand sandbox activity worker registration requires at least one registered activity.")
 
@@ -209,7 +209,7 @@ def _build_on_demand_sandbox_worker_start(
     return message
 
 
-def _build_on_demand_sandbox_worker_heartbeat(active_activities_count: int) -> pb.OnDemandSandboxActivityWorkerMessage:
+def build_on_demand_sandbox_worker_heartbeat(active_activities_count: int) -> pb.OnDemandSandboxActivityWorkerMessage:
     if active_activities_count < 0:
         raise ValueError("On-demand sandbox activity worker active activity count cannot be negative.")
 
@@ -223,7 +223,7 @@ def _normalize_optional_strings(values: Iterable[str]) -> list[str]:
 
 
 def _normalize_cpu(value: str) -> str:
-    normalized = _normalize_required(value, "On-demand sandbox activity declaration requires CPU resources.")
+    normalized = normalize_required(value, "On-demand sandbox activity declaration requires CPU resources.")
     milli_cpu = _try_parse_cpu_millicores(normalized)
     if milli_cpu is None or milli_cpu <= 0:
         raise ValueError(
@@ -233,7 +233,7 @@ def _normalize_cpu(value: str) -> str:
 
 
 def _normalize_memory(value: str) -> str:
-    normalized = _normalize_required(value, "On-demand sandbox activity declaration requires memory resources.")
+    normalized = normalize_required(value, "On-demand sandbox activity declaration requires memory resources.")
     memory_mib = _try_parse_memory_mib(normalized)
     if memory_mib is None or memory_mib <= 0:
         raise ValueError(

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

@@ -4,13 +4,13 @@
 from typing import Iterable, Optional
 
 
-def _normalize_required(value: Optional[str], message: str) -> str:
+def normalize_required(value: Optional[str], message: str) -> str:
     if not value or not value.strip():
         raise ValueError(message)
     return value.strip()
 
 
-def _resolve_activity_names(activity_names: str | Iterable[str]) -> list[str]:
+def resolve_activity_names(activity_names: str | Iterable[str]) -> list[str]:
     resolved: list[str] = []
     seen: set[str] = set()
     names = [activity_names] if isinstance(activity_names, str) else activity_names

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

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from typing import Iterable, Optional, Sequence
+from typing import Callable, Iterable, Optional, Protocol, Sequence, cast
 
 import grpc
 from azure.core.credentials import TokenCredential
@@ -15,6 +15,21 @@
 import durabletask.internal.shared as shared
 
 
+class _OnDemandSandboxActivitiesStub(Protocol):
+    DeclareOnDemandSandboxActivities: Callable[
+        [pb.OnDemandSandboxActivityDeclaration],
+        pb.OnDemandSandboxActivityDeclarationResult,
+    ]
+    RemoveOnDemandSandboxActivityDeclaration: Callable[
+        [pb.RemoveOnDemandSandboxActivityDeclarationRequest],
+        pb.RemoveOnDemandSandboxActivityDeclarationResult,
+    ]
+    ConnectOnDemandSandboxActivityWorker: Callable[
+        [Iterable[pb.OnDemandSandboxActivityWorkerMessage]],
+        pb.OnDemandSandboxActivityWorkerSessionResult,
+    ]
+
+
 class OnDemandSandboxActivitiesGrpcTransport:
     """Internal gRPC transport for on-demand sandbox activity RPCs."""
 
@@ -42,7 +57,7 @@ def __init__(
                 interceptors=resolved_interceptors,
                 channel_options=channel_options)
         self._channel = channel
-        self._stub = stubs.OnDemandSandboxActivitiesStub(channel)
+        self._stub = cast(_OnDemandSandboxActivitiesStub, stubs.OnDemandSandboxActivitiesStub(channel))
 
     def close(self) -> None:
         if self._owns_channel:

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

@@ -5,16 +5,19 @@
 import random
 import threading
 
-from typing import Iterator, Optional
+from typing import Any, Iterator, Optional
 
+from azure.core.credentials import TokenCredential
 from azure.identity import ManagedIdentityCredential
 
-from durabletask.azuremanaged.preview.on_demand_sandbox.helpers import _resolve_activity_names
+from durabletask import task
+from durabletask.azuremanaged.internal import on_demand_sandbox_activities_service_pb2 as pb
+from durabletask.azuremanaged.preview.on_demand_sandbox.helpers import resolve_activity_names
 from durabletask.azuremanaged.preview.on_demand_sandbox.declarations import (
     DEFAULT_MAX_CONCURRENT_ACTIVITIES,
     DEFAULT_WORKER_PROFILE_ID,
-    _build_on_demand_sandbox_worker_heartbeat,
-    _build_on_demand_sandbox_worker_start,
+    build_on_demand_sandbox_worker_heartbeat,
+    build_on_demand_sandbox_worker_start,
 )
 from durabletask.azuremanaged.preview.on_demand_sandbox.transport import (
     OnDemandSandboxActivitiesGrpcTransport,
@@ -35,7 +38,7 @@ class OnDemandSandboxWorker(DurableTaskSchedulerWorker):
     restricts dispatch to the activities registered on this worker.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         resolved_host_address = _resolve_host_address()
         resolved_taskhub = _resolve_taskhub()
         resolved_secure_channel = _resolve_secure_channel(resolved_host_address)
@@ -68,7 +71,7 @@ def __init__(self):
         self._on_demand_sandbox_active_activities = 0
         self._on_demand_sandbox_active_activities_lock = threading.Lock()
 
-    def add_activity(self, fn) -> str:
+    def add_activity(self, fn: task.Activity[Any, Any]) -> str:
         activity_name = super().add_activity(fn)
         self._on_demand_sandbox_activity_names.append(activity_name)
         return activity_name
@@ -82,16 +85,16 @@ def stop(self) -> None:
         self._stop_on_demand_sandbox_registration()
         super().stop()
 
-    def _durabletask_on_activity_execution_started(self, req) -> None:
+    def _durabletask_on_activity_execution_started(self, req: object) -> None:
         with self._on_demand_sandbox_active_activities_lock:
             self._on_demand_sandbox_active_activities += 1
 
-    def _durabletask_on_activity_execution_completed(self, req) -> None:
+    def _durabletask_on_activity_execution_completed(self, req: object) -> None:
         with self._on_demand_sandbox_active_activities_lock:
             self._on_demand_sandbox_active_activities = max(0, self._on_demand_sandbox_active_activities - 1)
 
     def _configure_on_demand_sandbox_activity_filters(self) -> None:
-        activity_names = _resolve_activity_names(self._on_demand_sandbox_activity_names)
+        activity_names = resolve_activity_names(self._on_demand_sandbox_activity_names)
         if not activity_names:
             raise RuntimeError(
                 "On-demand sandbox worker requires at least one registered activity before it can register.")
@@ -143,8 +146,8 @@ def _run_on_demand_sandbox_registration_loop(self) -> None:
                 self._on_demand_sandbox_registration_stop.wait(delay)
                 retry_delay = min(retry_delay * 2, 30.0)
 
-    def _registration_messages(self) -> Iterator:
-        yield _build_on_demand_sandbox_worker_start(
+    def _registration_messages(self) -> Iterator[pb.OnDemandSandboxActivityWorkerMessage]:
+        yield build_on_demand_sandbox_worker_start(
             taskhub=self._on_demand_sandbox_taskhub,
             worker_profile_id=self._on_demand_sandbox_worker_profile_id,
             max_activities_count=self._on_demand_sandbox_max_activities,
@@ -156,7 +159,7 @@ def _registration_messages(self) -> Iterator:
                 self._on_demand_sandbox_heartbeat_interval_seconds):
             with self._on_demand_sandbox_active_activities_lock:
                 active_count = self._on_demand_sandbox_active_activities
-            yield _build_on_demand_sandbox_worker_heartbeat(active_count)
+            yield build_on_demand_sandbox_worker_heartbeat(active_count)
 
 
 def _resolve_taskhub() -> str:
@@ -193,7 +196,7 @@ def _resolve_worker_profile_id() -> str:
     return resolved_worker_profile_id.strip()
 
 
-def _resolve_token_credential():
+def _resolve_token_credential() -> TokenCredential | None:
     authentication = os.getenv("DTS_AUTHENTICATION", "")
     if authentication.lower() != "managedidentity":
         return None