fix continue as new losing track of router and making placement fail

refactor: clean up debug print statements and improve code formatting

Removed unnecessary debug print statements from the TaskHubGrpcClient class to streamline the code and enhance readability. Additionally, adjusted formatting in various files for consistency, including property definitions and string formatting. This update aims to improve code maintainability and clarity.

Signed-off-by: Filinto Duran <[email protected]>

Author: filintod

durabletask/__init__.py

@@ -3,7 +3,6 @@
 
 """Durable Task SDK for Python"""
 
-
 PACKAGE_NAME = 'durabletask'
 
 # Public async exports (import directly from durabletask.aio)

durabletask/aio/compatibility.py

@@ -169,8 +169,8 @@ def ensure_compatibility(context_class: type) -> type:
 
     if missing_items:
         raise TypeError(
-            f"{context_class.__name__} does not implement OrchestrationContextProtocol. "
-            f"Missing: {', '.join(missing_items)}"
+            f'{context_class.__name__} does not implement OrchestrationContextProtocol. '
+            f'Missing: {", ".join(missing_items)}'
         )
 
     return context_class

durabletask/aio/context.py

@@ -78,6 +78,8 @@ class AsyncWorkflowContext(DeterministicContextMixin):
         '_detection_disabled',
         '_workflow_name',
         '_current_step',
+        '_sandbox_originals',
+        '_sandbox_mode',
     )
 
     # Generic type variable for awaitable result
@@ -95,6 +97,9 @@ def __init__(self, base_ctx: dt_task.OrchestrationContext):
         self._cleanup_tasks: list[Callable[[], Any]] = []
         self._workflow_name: Optional[str] = None
         self._current_step: Optional[str] = None
+        # Set by sandbox when active
+        self._sandbox_originals: Optional[dict[str, Any]] = None
+        self._sandbox_mode: Optional[str] = None
 
         # Performance optimization: Check if detection should be globally disabled
         self._detection_disabled = os.getenv('DAPR_WF_DISABLE_DETECTION') == 'true'
@@ -192,6 +197,37 @@ def info(self) -> WorkflowInfo:
             workflow_attempt=self.workflow_attempt,
         )
 
+    # Unsafe escape hatch: real wall-clock UTC now (best_effort only, not during replay)
+    def unsafe_wall_clock_now(self) -> datetime:
+        """
+        Return the real UTC wall-clock time.
+
+        Constraints:
+        - Raises RuntimeError if called during replay.
+        - Raises RuntimeError if sandbox mode is 'strict'.
+        - Intended for telemetry/logging only; do not use for workflow decisions.
+        """
+        if self.is_replaying:
+            raise RuntimeError('unsafe_wall_clock_now() cannot be used during replay')
+        mode = getattr(self, '_sandbox_mode', None)
+        if mode == 'strict':
+            raise RuntimeError('unsafe_wall_clock_now() is disabled in strict sandbox mode')
+        originals = getattr(self, '_sandbox_originals', None)
+        if not originals or 'time.time' not in originals:
+            # Fallback to system if sandbox not active
+            import time as _time
+            from datetime import timezone
+
+            return datetime.fromtimestamp(_time.time(), tz=timezone.utc)
+        real_time = originals['time.time']
+        try:
+            from datetime import timezone
+
+            ts = float(real_time())  # type: ignore[call-arg]
+            return datetime.fromtimestamp(ts, tz=timezone.utc)
+        except Exception as e:
+            raise RuntimeError(f'unsafe_wall_clock_now() failed: {e}')
+
     # Activity operations
     def activity(
         self,

durabletask/aio/driver.py

@@ -35,8 +35,7 @@
 class WorkflowFunction(Protocol):
     """Protocol for workflow functions."""
 
-    async def __call__(self, ctx: Any, input_data: Optional[Any] = None) -> Any:
-        ...
+    async def __call__(self, ctx: Any, input_data: Optional[Any] = None) -> Any: ...
 
 
 class CoroutineOrchestratorRunner:

durabletask/aio/errors.py

@@ -45,7 +45,7 @@ def __init__(
         if step:
             context_parts.append(f'step: {step}')
 
-        context_str = f" ({', '.join(context_parts)})" if context_parts else ''
+        context_str = f' ({", ".join(context_parts)})' if context_parts else ''
         super().__init__(f'{message}{context_str}')
 
 

durabletask/aio/sandbox.py

@@ -293,6 +293,14 @@ def __enter__(self) -> '_Sandbox':
         # Apply patches for best_effort and strict modes
         self._apply_patches()
 
+        # Expose originals/mode to the async workflow context for controlled unsafe access
+        try:
+            setattr(self.async_ctx, '_sandbox_originals', dict(self.originals))
+            setattr(self.async_ctx, '_sandbox_mode', self.mode)
+        except Exception:
+            # Context may not support attribute assignment; ignore
+            pass
+
         return self
 
     def __exit__(
@@ -307,6 +315,15 @@ def __exit__(
         if self.mode != 'off' and self.originals:
             self._restore_originals()
 
+        # Remove exposed references from the async context
+        try:
+            if hasattr(self.async_ctx, '_sandbox_originals'):
+                delattr(self.async_ctx, '_sandbox_originals')
+            if hasattr(self.async_ctx, '_sandbox_mode'):
+                delattr(self.async_ctx, '_sandbox_mode')
+        except Exception:
+            pass
+
     def _apply_patches(self) -> None:
         """Apply patches to non-deterministic functions."""
         import asyncio as _asyncio
@@ -525,9 +542,7 @@ async def _empty() -> list[Any]:
             if all(_is_workflow_awaitable(a) for a in aws):
 
                 async def _await_when_all() -> Any:
-                    from .awaitables import (
-                        WhenAllAwaitable,  # local import to avoid cycles
-                    )
+                    from .awaitables import WhenAllAwaitable  # local import to avoid cycles
 
                     combined: Any = WhenAllAwaitable(list(aws))
                     return await combined

durabletask/client.py

@@ -232,13 +232,9 @@ def wait_for_orchestration_completion(
             # For positive timeout, best-effort pre-check and short polling to avoid long server waits
             try:
                 # First check if the orchestration is already completed
-                print(f'[CLIENT DEBUG] Checking current state before waiting for {instance_id}')
                 current_state = self.get_orchestration_state(
                     instance_id, fetch_payloads=fetch_payloads
                 )
-                print(
-                    f"[CLIENT DEBUG] Current state: {current_state.runtime_status if current_state else 'None'}"
-                )
                 if current_state and current_state.runtime_status in [
                     OrchestrationStatus.COMPLETED,
                     OrchestrationStatus.FAILED,
@@ -256,42 +252,28 @@ def wait_for_orchestration_completion(
                 poll_start = time.time()
                 poll_interval = 0.1
 
-                print(f'[CLIENT DEBUG] Starting polling for up to {poll_timeout}s')
                 while time.time() - poll_start < poll_timeout:
                     current_state = self.get_orchestration_state(
                         instance_id, fetch_payloads=fetch_payloads
                     )
-                    print(
-                        f"[CLIENT DEBUG] Poll state: {current_state.runtime_status if current_state else 'None'}"
-                    )
 
                     if current_state and current_state.runtime_status in [
                         OrchestrationStatus.COMPLETED,
                         OrchestrationStatus.FAILED,
                         OrchestrationStatus.TERMINATED,
                     ]:
-                        print(
-                            f'[CLIENT DEBUG] Found terminal state during polling: {current_state.runtime_status}'
-                        )
                         return current_state
 
                     time.sleep(poll_interval)
                     poll_interval = min(poll_interval * 1.5, 1.0)  # Exponential backoff, max 1s
-                print(f'[CLIENT DEBUG] Polling completed, no terminal state found')
             except Exception:
                 # Ignore pre-check/poll issues (e.g., mocked stubs in unit tests) and fall back
                 pass
 
             self._logger.info(f"Waiting up to {timeout}s for instance '{instance_id}' to complete.")
-            print(
-                f'[CLIENT DEBUG] About to call WaitForInstanceCompletion for {instance_id} with timeout {grpc_timeout}'
-            )
             res: pb.GetInstanceResponse = self._stub.WaitForInstanceCompletion(
                 req, timeout=grpc_timeout
             )
-            print(
-                f'[CLIENT DEBUG] WaitForInstanceCompletion returned successfully for {instance_id}'
-            )
             state = new_orchestration_state(req.instanceId, res)
             if not state:
                 return None
@@ -311,17 +293,10 @@ def wait_for_orchestration_completion(
 
             return state
         except grpc.RpcError as rpc_error:
-            print(
-                f'[CLIENT DEBUG] gRPC error in WaitForInstanceCompletion for {instance_id}: {rpc_error.code()} - {rpc_error.details()}'
-            )
             if rpc_error.code() == grpc.StatusCode.DEADLINE_EXCEEDED:  # type: ignore
-                print(
-                    f'[CLIENT DEBUG] Deadline exceeded for {instance_id}, converting to TimeoutError'
-                )
                 # Replace gRPC error with the built-in TimeoutError
                 raise TimeoutError('Timed-out waiting for the orchestration to complete')
             else:
-                print(f'[CLIENT DEBUG] Re-raising non-timeout gRPC error for {instance_id}')
                 raise
 
     def raise_orchestration_event(

durabletask/deterministic.py

@@ -64,12 +64,10 @@ class DeterministicContextProtocol(Protocol):
     """Protocol for contexts that provide deterministic operations."""
 
     @property
-    def instance_id(self) -> str:
-        ...
+    def instance_id(self) -> str: ...
 
     @property
-    def current_utc_datetime(self) -> datetime:
-        ...
+    def current_utc_datetime(self) -> datetime: ...
 
 
 class DeterministicContextMixin:

durabletask/internal/helpers.py

@@ -208,6 +208,7 @@ def new_complete_orchestration_action(
     result: Optional[str] = None,
     failure_details: Optional[pb.TaskFailureDetails] = None,
     carryover_events: Optional[list[pb.HistoryEvent]] = None,
+    app_id: Optional[str] = None,
 ) -> pb.OrchestratorAction:
     completeOrchestrationAction = pb.CompleteOrchestrationAction(
         orchestrationStatus=status,
@@ -216,7 +217,10 @@ def new_complete_orchestration_action(
         carryoverEvents=carryover_events,
     )
 
-    return pb.OrchestratorAction(id=id, completeOrchestration=completeOrchestrationAction)
+    return pb.OrchestratorAction(
+        id=id, completeOrchestration=completeOrchestrationAction, 
+        router=pb.TaskRouter(sourceAppID=app_id) if app_id else None
+    )
 
 
 def new_create_timer_action(id: int, fire_at: datetime) -> pb.OrchestratorAction:

durabletask/internal/shared.py

@@ -181,8 +181,8 @@ def get_grpc_retry_service_config_option() -> Optional[tuple[str, str]]:
                 'name': [{'service': ''}],
                 'retryPolicy': {
                     'maxAttempts': max_attempts,
-                    'initialBackoff': f'{initial_backoff_ms/1000.0}s',
-                    'maxBackoff': f'{max_backoff_ms/1000.0}s',
+                    'initialBackoff': f'{initial_backoff_ms / 1000.0}s',
+                    'maxBackoff': f'{max_backoff_ms / 1000.0}s',
                     'backoffMultiplier': backoff_multiplier,
                     'retryableStatusCodes': codes,
                 },
@@ -222,7 +222,7 @@ def get_logger(
     # Add a default log handler if none is provided
     if log_handler is None:
         log_handler = logging.StreamHandler()
-        log_handler.setLevel(logging.INFO)
+        log_handler.setLevel(logging.DEBUG)
     logger.handlers.append(log_handler)
 
     # Set a default log formatter to our handler if none is provided