wire deterministic to main running context, increase coverage on worker tests, add reduce stop timeout for task runner

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

Author: filintod

Makefile

@@ -7,6 +7,13 @@ test-unit:
 test-e2e:
 	pytest -m e2e --verbose
 
+coverage-clean:
+	rm -f .coverage .coverage.* coverage.xml
+
+coverage-all: coverage-clean
+	pytest -m "not e2e" --cov=durabletask --cov-branch --cov-report=term-missing --cov-report=xml
+	pytest -m e2e --cov=durabletask --cov-branch --cov-report=term-missing --cov-report=xml --cov-append
+
 install:
 	python3 -m pip install .
 
@@ -18,4 +25,4 @@ gen-proto:
 	python3 -m grpc_tools.protoc --proto_path=.  --python_out=. --pyi_out=. --grpc_python_out=. ./durabletask/internal/orchestrator_service.proto
 	rm durabletask/internal/*.proto
 
-.PHONY: init test-unit test-e2e gen-proto install
+.PHONY: init test-unit test-e2e coverage-clean coverage-unit coverage-e2e coverage-all gen-proto install

durabletask/deterministic.py

@@ -25,8 +25,8 @@
 import uuid
 from collections.abc import Sequence
 from dataclasses import dataclass
-from datetime import datetime
-from typing import Optional, Protocol, TypeVar, runtime_checkable
+from datetime import datetime, timedelta
+from typing import Optional, TypeVar
 
 
 @dataclass
@@ -99,17 +99,6 @@ def deterministic_uuid_v5(instance_id: str, current_datetime: datetime, counter:
     return uuid.uuid5(namespace, name)
 
 
-@runtime_checkable
-class DeterministicContextProtocol(Protocol):
-    """Protocol for contexts that provide deterministic operations."""
-
-    @property
-    def instance_id(self) -> str: ...
-
-    @property
-    def current_utc_datetime(self) -> datetime: ...
-
-
 class DeterministicContextMixin:
     """
     Mixin providing deterministic helpers for workflow contexts.
@@ -121,25 +110,25 @@ class DeterministicContextMixin:
     """
 
     def __init__(self, *args, **kwargs):
-        """Initialize the mixin with a UUID counter."""
+        """Initialize the mixin with UUID and timestamp counters."""
         super().__init__(*args, **kwargs)
         # Counter for deterministic UUID generation (matches .NET newGuidCounter)
         # This counter resets to 0 on each replay, ensuring determinism
         self._uuid_counter: int = 0
+        # Counter for deterministic timestamp sequencing (resets on replay)
+        self._timestamp_counter: int = 0
 
     def now(self) -> datetime:
-        """Return orchestration time (deterministic UTC)."""
-        value = self.current_utc_datetime  # type: ignore[attr-defined]
-        assert isinstance(value, datetime)
-        return value
+        """Alias for deterministic current_utc_datetime."""
+        return self.current_utc_datetime  # type: ignore[attr-defined]
 
     def random(self) -> random.Random:
         """Return a PRNG seeded deterministically from instance id and orchestration time."""
         rnd = deterministic_random(
             self.instance_id,  # type: ignore[attr-defined]
             self.current_utc_datetime,  # type: ignore[attr-defined]
         )
-        # Mark as deterministic for sandbox detector whitelisting of bound methods
+        # Mark as deterministic for asyncio sandbox detector whitelisting of bound methods (randint, random)
         try:
             rnd._dt_deterministic = True
         except Exception:
@@ -201,3 +190,35 @@ def random_choice(self, sequence: Sequence[T]) -> T:
             raise IndexError("Cannot choose from empty sequence")
         rnd = self.random()
         return rnd.choice(sequence)
+
+    def now_with_sequence(self) -> datetime:
+        """
+        Return deterministic timestamp with microsecond increment per call.
+
+        Each call returns: current_utc_datetime + (counter * 1 microsecond)
+
+        This provides ordered, unique timestamps for tracing/telemetry while maintaining
+        determinism across replays. The counter resets to 0 on each replay (similar to
+        _uuid_counter pattern).
+
+        Perfect for preserving event ordering within a workflow without requiring activities.
+
+        Returns:
+            datetime: Deterministic timestamp that increments on each call
+
+        Example:
+            ```python
+            def workflow(ctx):
+                t1 = ctx.now_with_sequence()  # 2024-01-01 12:00:00.000000
+                result = yield ctx.call_activity(some_activity, input="data")
+                t2 = ctx.now_with_sequence()  # 2024-01-01 12:00:00.000001
+                # t1 < t2, preserving order for telemetry
+            ```
+        """
+        offset = timedelta(microseconds=self._timestamp_counter)
+        self._timestamp_counter += 1
+        return self.current_utc_datetime + offset  # type: ignore[attr-defined]
+
+    def current_utc_datetime_with_sequence(self):
+        """Alias for now_with_sequence for API parity with other SDKs."""
+        return self.now_with_sequence()

durabletask/worker.py

@@ -19,7 +19,7 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
-from durabletask import task
+from durabletask import deterministic, task
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
 
 TInput = TypeVar("TInput")
@@ -159,6 +159,8 @@ class TaskHubGrpcWorker:
             interceptors to apply to the channel. Defaults to None.
         concurrency_options (Optional[ConcurrencyOptions], optional): Configuration for
             controlling worker concurrency limits. If None, default settings are used.
+        stop_timeout (float, optional): Maximum time in seconds to wait for the worker thread
+            to stop when calling stop(). Defaults to 30.0. Useful to set lower values in tests.
 
     Attributes:
         concurrency_options (ConcurrencyOptions): The current concurrency configuration.
@@ -224,6 +226,7 @@ def __init__(
         interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
         concurrency_options: Optional[ConcurrencyOptions] = None,
         channel_options: Optional[Sequence[tuple[str, Any]]] = None,
+        stop_timeout: float = 30.0,
     ):
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
@@ -232,6 +235,7 @@ def __init__(
         self._is_running = False
         self._secure_channel = secure_channel
         self._channel_options = channel_options
+        self._stop_timeout = stop_timeout
         # Track in-flight activity executions for graceful draining
         import threading as _threading
 
@@ -512,7 +516,7 @@ def stop(self):
         if self._response_stream is not None:
             self._response_stream.cancel()
         if self._runLoop is not None:
-            self._runLoop.join(timeout=30)
+            self._runLoop.join(timeout=self._stop_timeout)
         self._async_worker_manager.shutdown()
         self._logger.info("Worker shutdown completed")
         self._is_running = False
@@ -659,11 +663,12 @@ def _execute_activity(
             )
 
 
-class _RuntimeOrchestrationContext(task.OrchestrationContext):
+class _RuntimeOrchestrationContext(task.OrchestrationContext, deterministic.DeterministicContextMixin):
     _generator: Optional[Generator[task.Task, Any, Any]]
     _previous_task: Optional[task.Task]
 
     def __init__(self, instance_id: str):
+        super().__init__()
         self._generator = None
         self._is_replaying = True
         self._is_suspended = False