Merge branch 'main' into claude-1

Author: timl3136

.github/workflows/ci_checks.yml

@@ -51,10 +51,10 @@ jobs:
         uses: astral-sh/setup-uv@v1
       - name: Install dependencies
         run: |
-          uv sync --extra dev
+          uv sync --extra dev --extra openai
       - name: Run mypy type checker
         run: |
-          uv tool run mypy cadence/
+          uv run mypy cadence/
 
   test:
     name: Unit Tests

cadence/_internal/activity/_activity_executor.py

@@ -1,12 +1,13 @@
 from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
 from traceback import format_exception
-from typing import Any, Callable, cast
+from typing import Any, Callable, Union, cast
 from google.protobuf.duration import to_timedelta
 from google.protobuf.timestamp import to_datetime
 
 from cadence._internal.activity._context import _Context, _SyncContext
 from cadence._internal.activity._definition import BaseDefinition, ExecutionStrategy
+from cadence._internal.activity._heartbeat import _HeartbeatSender
 from cadence.activity import ActivityInfo, ActivityDefinition
 from cadence.api.v1.common_pb2 import Failure
 from cadence.api.v1.service_worker_pb2 import (
@@ -46,19 +47,33 @@ async def execute(self, task: PollForActivityTaskResponse):
             _logger.exception("Activity failed")
             await self._report_failure(task, e)
 
-    def _create_context(self, task: PollForActivityTaskResponse) -> _Context:
+    def _create_context(
+        self, task: PollForActivityTaskResponse
+    ) -> Union[_Context, _SyncContext]:
         activity_type = task.activity_type.name
         try:
             activity_def = cast(BaseDefinition, self._registry(activity_type))
         except KeyError:
             raise KeyError(f"Activity type not found: {activity_type}") from None
 
         info = self._create_info(task)
+        heartbeat_sender = _HeartbeatSender(
+            self._client.worker_stub,
+            self._data_converter,
+            task.task_token,
+            self._identity,
+        )
 
         if activity_def.strategy == ExecutionStrategy.ASYNC:
-            return _Context(self._client, info, activity_def)
+            return _Context(self._client, info, activity_def, heartbeat_sender)
         else:
-            return _SyncContext(self._client, info, activity_def, self._thread_pool)
+            return _SyncContext(
+                self._client,
+                info,
+                activity_def,
+                self._thread_pool,
+                heartbeat_sender,
+            )
 
     async def _report_failure(
         self, task: PollForActivityTaskResponse, error: Exception

cadence/_internal/activity/_context.py

@@ -4,6 +4,7 @@
 
 from cadence import Client
 from cadence._internal.activity._definition import BaseDefinition
+from cadence._internal.activity._heartbeat import _HeartbeatSender
 from cadence.activity import ActivityInfo, ActivityContext
 from cadence.api.v1.common_pb2 import Payload
 
@@ -14,15 +15,27 @@ def __init__(
         client: Client,
         info: ActivityInfo,
         activity_def: BaseDefinition[[Any], Any],
+        heartbeat_sender: _HeartbeatSender,
     ):
         self._client = client
         self._info = info
         self._activity_def = activity_def
+        self._heartbeat_sender = heartbeat_sender
+        self._heartbeat_tasks: set[asyncio.Future[None]] = set()
 
     async def execute(self, payload: Payload) -> Any:
         params = self._to_params(payload)
-        with self._activate():
-            return await self._activity_def.impl_fn(*params)
+        try:
+            with self._activate():
+                return await self._activity_def.impl_fn(*params)
+        finally:
+            await self._wait_pending_heartbeats()
+
+    async def _wait_pending_heartbeats(self) -> None:
+        if not self._heartbeat_tasks:
+            return
+        tasks = list(self._heartbeat_tasks)
+        await asyncio.gather(*tasks, return_exceptions=True)
 
     def _to_params(self, payload: Payload) -> list[Any]:
         return self._activity_def.signature.params_from_payload(
@@ -35,6 +48,13 @@ def client(self) -> Client:
     def info(self) -> ActivityInfo:
         return self._info
 
+    def heartbeat(self, *details: Any) -> None:
+        heartbeat_task = asyncio.create_task(
+            self._heartbeat_sender.send_heartbeat(*details)
+        )
+        self._heartbeat_tasks.add(heartbeat_task)
+        heartbeat_task.add_done_callback(self._heartbeat_tasks.discard)
+
 
 class _SyncContext(_Context):
     def __init__(
@@ -43,18 +63,30 @@ def __init__(
         info: ActivityInfo,
         activity_def: BaseDefinition[[Any], Any],
         executor: ThreadPoolExecutor,
+        heartbeat_sender: _HeartbeatSender,
     ):
-        super().__init__(client, info, activity_def)
+        super().__init__(client, info, activity_def, heartbeat_sender)
         self._executor = executor
 
     async def execute(self, payload: Payload) -> Any:
         params = self._to_params(payload)
-        loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(self._executor, self._run, params)
+        self._loop = asyncio.get_running_loop()
+        try:
+            return await self._loop.run_in_executor(self._executor, self._run, params)
+        finally:
+            await self._wait_pending_heartbeats()
 
     def _run(self, args: list[Any]) -> Any:
         with self._activate():
             return self._activity_def.impl_fn(*args)
 
     def client(self) -> Client:
         raise RuntimeError("client is only supported in async activities")
+
+    def heartbeat(self, *details: Any) -> None:
+        future = asyncio.run_coroutine_threadsafe(
+            self._heartbeat_sender.send_heartbeat(*details), self._loop
+        )
+        wrapped = asyncio.wrap_future(future, loop=self._loop)
+        self._heartbeat_tasks.add(wrapped)
+        wrapped.add_done_callback(self._heartbeat_tasks.discard)

cadence/_internal/activity/_definition.py

@@ -14,7 +14,9 @@
     ParamSpec,
     TypeVar,
     Awaitable,
+    Type,
     cast,
+    overload,
     Concatenate,
 )
 
@@ -119,10 +121,13 @@ def __init__(
         super().__init__(name, wrapped, ExecutionStrategy.THREAD_POOL, signature)
         update_wrapper(self, wrapped)
 
-    def __get__(self, instance, owner):
+    @overload
+    def __get__(self, instance: None, owner: Type[T]) -> "SyncMethodImpl[T, P, R]": ...
+    @overload
+    def __get__(self, instance: T, owner: Type[T]) -> SyncImpl[P, R]: ...
+    def __get__(self, instance: T | None, owner: Type[T]) -> "SyncImpl[P, R] | Self":
         if instance is None:
             return self
-        # If we bound the method to an instance, then drop the self parameter. It's a normal function again
         return SyncImpl[P, R](
             partial(self._wrapped, instance), self.name, self._signature
         )
@@ -181,10 +186,13 @@ def __init__(
         else:
             self._is_coroutine = _COROUTINE_MARKER
 
-    def __get__(self, instance, owner):
+    @overload
+    def __get__(self, instance: None, owner: Type[T]) -> "AsyncMethodImpl[T, P, R]": ...
+    @overload
+    def __get__(self, instance: T, owner: Type[T]) -> AsyncImpl[P, R]: ...
+    def __get__(self, instance: T | None, owner: Type[T]) -> "AsyncImpl[P, R] | Self":
         if instance is None:
             return self
-        # If we bound the method to an instance, then drop the self parameter. It's a normal function again
         return AsyncImpl[P, R](
             partial(self._wrapped, instance), self.name, self._signature
         )

cadence/_internal/activity/_heartbeat.py

@@ -0,0 +1,35 @@
+from logging import getLogger
+from typing import Any
+
+from cadence.api.v1.service_worker_pb2 import RecordActivityTaskHeartbeatRequest
+from cadence.api.v1.service_worker_pb2_grpc import WorkerAPIStub
+from cadence.data_converter import DataConverter
+
+_logger = getLogger(__name__)
+
+
+class _HeartbeatSender:
+    def __init__(
+        self,
+        worker_stub: WorkerAPIStub,
+        data_converter: DataConverter,
+        task_token: bytes,
+        identity: str,
+    ):
+        self._worker_stub = worker_stub
+        self._data_converter = data_converter
+        self._task_token = task_token
+        self._identity = identity
+
+    async def send_heartbeat(self, *details: Any) -> None:
+        try:
+            payload = self._data_converter.to_data(list(details))
+            await self._worker_stub.RecordActivityTaskHeartbeat(
+                RecordActivityTaskHeartbeatRequest(
+                    task_token=self._task_token,
+                    details=payload,
+                    identity=self._identity,
+                )
+            )
+        except Exception:
+            _logger.warning("Heartbeat failed", exc_info=True)

cadence/_internal/workflow/workflow_engine.py

@@ -13,16 +13,19 @@
 from cadence._internal.workflow.statemachine.decision_manager import DecisionManager
 from cadence._internal.workflow.workflow_instance import WorkflowInstance
 from cadence.api.v1 import history
-from cadence.api.v1.common_pb2 import Failure
+from cadence.api.v1.common_pb2 import Failure, WorkflowType
 from cadence.api.v1.decision_pb2 import (
     Decision,
     FailWorkflowExecutionDecisionAttributes,
     CompleteWorkflowExecutionDecisionAttributes,
+    ContinueAsNewWorkflowExecutionDecisionAttributes,
 )
 from cadence.api.v1.history_pb2 import (
     HistoryEvent,
     WorkflowExecutionStartedEventAttributes,
 )
+from cadence.api.v1.tasklist_pb2 import TaskList
+from cadence.error import ContinueAsNewError
 from cadence.workflow import WorkflowDefinition, WorkflowInfo
 
 logger = logging.getLogger(__name__)
@@ -187,6 +190,25 @@ def _maybe_complete_workflow(self) -> Optional[Decision]:
                 return None
         except (CancelledError, InvalidStateError, FatalDecisionError):
             raise
+        except ContinueAsNewError as e:
+            # Use execution's workflow type and task list when not overridden
+            info = self._context.info()
+            attrs = ContinueAsNewWorkflowExecutionDecisionAttributes(
+                workflow_type=WorkflowType(name=e.workflow_type or info.workflow_type),
+                task_list=TaskList(name=e.task_list or info.workflow_task_list),
+                input=info.data_converter.to_data(list(e.workflow_args)),
+            )
+            if e.execution_start_to_close_timeout is not None:
+                attrs.execution_start_to_close_timeout.FromTimedelta(
+                    e.execution_start_to_close_timeout
+                )
+            if e.task_start_to_close_timeout is not None:
+                attrs.task_start_to_close_timeout.FromTimedelta(
+                    e.task_start_to_close_timeout
+                )
+            return Decision(
+                continue_as_new_workflow_execution_decision_attributes=attrs,
+            )
         except ExceptionGroup as e:
             if e.subgroup((InvalidStateError, FatalDecisionError)):
                 raise

cadence/activity.py

@@ -63,6 +63,11 @@ def info() -> ActivityInfo:
     return ActivityContext.get().info()
 
 
+def heartbeat(*details: Any) -> None:
+    """Send a heartbeat for the current activity."""
+    ActivityContext.get().heartbeat(*details)
+
+
 class ActivityContext(ABC):
     _var: ContextVar["ActivityContext"] = ContextVar("activity")
 
@@ -72,6 +77,9 @@ def info(self) -> ActivityInfo: ...
     @abstractmethod
     def client(self) -> Client: ...
 
+    @abstractmethod
+    def heartbeat(self, *details: Any) -> None: ...
+
     @contextmanager
     def _activate(self) -> Iterator[None]:
         token = ActivityContext._var.set(self)