hatchet-sdk 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

hatchet_sdk/clients/admin.py

@@ -5,18 +5,27 @@ from typing import Union, cast
 
 import grpc
 from google.protobuf import timestamp_pb2
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
 from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
-from hatchet_sdk.contracts import workflows_pb2 as workflow_protos
+from hatchet_sdk.contracts import workflows_pb2 as v0_workflow_protos
+from hatchet_sdk.contracts.v1 import workflows_pb2 as workflow_protos
+from hatchet_sdk.contracts.v1.workflows_pb2_grpc import AdminServiceStub
 from hatchet_sdk.contracts.workflows_pb2_grpc import WorkflowServiceStub
 from hatchet_sdk.metadata import get_metadata
 from hatchet_sdk.rate_limit import RateLimitDuration
-from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto, maybe_int_to_str
+from hatchet_sdk.runnables.contextvars import (
+    ctx_step_run_id,
+    ctx_worker_id,
+    ctx_workflow_run_id,
+    spawn_index_lock,
+    workflow_spawn_indices,
+)
+from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 from hatchet_sdk.workflow_run import WorkflowRunRef
 
@@ -29,28 +38,19 @@ class ScheduleTriggerWorkflowOptions(BaseModel):
     namespace: str | None = None
 
 
-class ChildTriggerWorkflowOptions(BaseModel):
-    additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
-    sticky: bool = False
-
-
-class ChildWorkflowRunDict(BaseModel):
-    workflow_name: str
-    input: JSONSerializableMapping
-    options: ChildTriggerWorkflowOptions
-    key: str | None = None
-
-
 class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions):
     additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
     desired_worker_id: str | None = None
     namespace: str | None = None
+    sticky: bool = False
+    key: str | None = None
 
 
-class WorkflowRunDict(BaseModel):
+class WorkflowRunTriggerConfig(BaseModel):
     workflow_name: str
     input: JSONSerializableMapping
     options: TriggerWorkflowOptions
+    key: str | None = None
 
 
 class DedupeViolationErr(Exception):
@@ -63,55 +63,69 @@ class AdminClient:
     def __init__(self, config: ClientConfig):
         conn = new_conn(config, False)
         self.config = config
-        self.client = WorkflowServiceStub(conn)  # type: ignore[no-untyped-call]
+        self.client = AdminServiceStub(conn)  # type: ignore[no-untyped-call]
+        self.v0_client = WorkflowServiceStub(conn)  # type: ignore[no-untyped-call]
         self.token = config.token
         self.listener_client = RunEventListenerClient(config=config)
         self.namespace = config.namespace
 
         self.pooled_workflow_listener: PooledWorkflowRunListener | None = None
 
+    class TriggerWorkflowRequest(BaseModel):
+        model_config = ConfigDict(extra="ignore")
+
+        parent_id: str | None = None
+        parent_step_run_id: str | None = None
+        child_index: int | None = None
+        child_key: str | None = None
+        additional_metadata: str | None = None
+        desired_worker_id: str | None = None
+        priority: int | None = None
+
+        @field_validator("additional_metadata", mode="before")
+        @classmethod
+        def validate_additional_metadata(
+            cls, v: JSONSerializableMapping | None
+        ) -> bytes | None:
+            if not v:
+                return None
+
+            try:
+                return json.dumps(v).encode("utf-8")
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Error encoding payload: {e}")
+
     def _prepare_workflow_request(
         self,
         workflow_name: str,
         input: JSONSerializableMapping,
         options: TriggerWorkflowOptions,
-    ) -> workflow_protos.TriggerWorkflowRequest:
+    ) -> v0_workflow_protos.TriggerWorkflowRequest:
         try:
             payload_data = json.dumps(input)
-            _options = options.model_dump()
-
-            _options.pop("namespace")
-
-            try:
-                _options = {
-                    **_options,
-                    "additional_metadata": json.dumps(
-                        options.additional_metadata
-                    ).encode("utf-8"),
-                }
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Error encoding payload: {e}")
-
-            return workflow_protos.TriggerWorkflowRequest(
-                name=workflow_name, input=payload_data, **_options
-            )
         except json.JSONDecodeError as e:
             raise ValueError(f"Error encoding payload: {e}")
 
+        _options = self.TriggerWorkflowRequest.model_validate(
+            options.model_dump()
+        ).model_dump()
+
+        return v0_workflow_protos.TriggerWorkflowRequest(
+            name=workflow_name, input=payload_data, **_options
+        )
+
     def _prepare_put_workflow_request(
         self,
         name: str,
-        workflow: workflow_protos.CreateWorkflowVersionOpts,
-        overrides: workflow_protos.CreateWorkflowVersionOpts | None = None,
-    ) -> workflow_protos.PutWorkflowRequest:
+        workflow: workflow_protos.CreateWorkflowVersionRequest,
+        overrides: workflow_protos.CreateWorkflowVersionRequest | None = None,
+    ) -> workflow_protos.CreateWorkflowVersionRequest:
         if overrides is not None:
             workflow.MergeFrom(overrides)
 
         workflow.name = name
 
-        return workflow_protos.PutWorkflowRequest(
-            opts=workflow,
-        )
+        return workflow
 
     def _parse_schedule(
         self, schedule: datetime | timestamp_pb2.Timestamp
@@ -134,57 +148,21 @@ class AdminClient:
         schedules: list[Union[datetime, timestamp_pb2.Timestamp]],
         input: JSONSerializableMapping = {},
         options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
-    ) -> workflow_protos.ScheduleWorkflowRequest:
-        return workflow_protos.ScheduleWorkflowRequest(
+    ) -> v0_workflow_protos.ScheduleWorkflowRequest:
+        return v0_workflow_protos.ScheduleWorkflowRequest(
             name=name,
             schedules=[self._parse_schedule(schedule) for schedule in schedules],
             input=json.dumps(input),
             **options.model_dump(),
         )
 
-    ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
-    @tenacity_retry
-    async def aio_run_workflow(
-        self,
-        workflow_name: str,
-        input: JSONSerializableMapping,
-        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
-    ) -> WorkflowRunRef:
-        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
-        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
-        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
-        return await asyncio.to_thread(self.run_workflow, workflow_name, input, options)
-
-    @tenacity_retry
-    async def aio_run_workflows(
-        self,
-        workflows: list[WorkflowRunDict],
-        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
-    ) -> list[WorkflowRunRef]:
-        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
-        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
-        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
-        return await asyncio.to_thread(self.run_workflows, workflows, options)
-
     @tenacity_retry
     async def aio_put_workflow(
         self,
         name: str,
-        workflow: workflow_protos.CreateWorkflowVersionOpts,
-        overrides: workflow_protos.CreateWorkflowVersionOpts | None = None,
-    ) -> workflow_protos.WorkflowVersion:
-        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
-        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
-        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
+        workflow: workflow_protos.CreateWorkflowVersionRequest,
+        overrides: workflow_protos.CreateWorkflowVersionRequest | None = None,
+    ) -> workflow_protos.CreateWorkflowVersionResponse:
         return await asyncio.to_thread(self.put_workflow, name, workflow, overrides)
 
     @tenacity_retry
@@ -194,12 +172,6 @@ class AdminClient:
         limit: int,
         duration: RateLimitDuration = RateLimitDuration.SECOND,
     ) -> None:
-        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
-        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
-        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
         return await asyncio.to_thread(self.put_rate_limit, key, limit, duration)
 
     @tenacity_retry
@@ -209,13 +181,7 @@ class AdminClient:
         schedules: list[Union[datetime, timestamp_pb2.Timestamp]],
         input: JSONSerializableMapping = {},
         options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
-    ) -> workflow_protos.WorkflowVersion:
-        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
-        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
-        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
-        if not self.pooled_workflow_listener:
-            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
-
+    ) -> v0_workflow_protos.WorkflowVersion:
         return await asyncio.to_thread(
             self.schedule_workflow, name, schedules, input, options
         )
@@ -224,18 +190,19 @@ class AdminClient:
     def put_workflow(
         self,
         name: str,
-        workflow: workflow_protos.CreateWorkflowVersionOpts,
-        overrides: workflow_protos.CreateWorkflowVersionOpts | None = None,
-    ) -> workflow_protos.WorkflowVersion:
+        workflow: workflow_protos.CreateWorkflowVersionRequest,
+        overrides: workflow_protos.CreateWorkflowVersionRequest | None = None,
+    ) -> workflow_protos.CreateWorkflowVersionResponse:
         opts = self._prepare_put_workflow_request(name, workflow, overrides)
 
-        resp: workflow_protos.WorkflowVersion = self.client.PutWorkflow(
-            opts,
-            metadata=get_metadata(self.token),
+        return cast(
+            workflow_protos.CreateWorkflowVersionResponse,
+            self.client.PutWorkflow(
+                opts,
+                metadata=get_metadata(self.token),
+            ),
         )
 
-        return resp
-
     @tenacity_retry
     def put_rate_limit(
         self,
@@ -246,11 +213,12 @@ class AdminClient:
         duration_proto = convert_python_enum_to_proto(
             duration, workflow_protos.RateLimitDuration
         )
-        self.client.PutRateLimit(
-            workflow_protos.PutRateLimitRequest(
+
+        self.v0_client.PutRateLimit(
+            v0_workflow_protos.PutRateLimitRequest(
                 key=key,
                 limit=limit,
-                duration=maybe_int_to_str(duration_proto),
+                duration=duration_proto,  # type: ignore[arg-type]
             ),
             metadata=get_metadata(self.token),
         )
@@ -262,7 +230,7 @@ class AdminClient:
         schedules: list[Union[datetime, timestamp_pb2.Timestamp]],
         input: JSONSerializableMapping = {},
         options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
-    ) -> workflow_protos.WorkflowVersion:
+    ) -> v0_workflow_protos.WorkflowVersion:
         try:
             namespace = options.namespace or self.namespace
 
@@ -274,8 +242,8 @@ class AdminClient:
             )
 
             return cast(
-                workflow_protos.WorkflowVersion,
-                self.client.ScheduleWorkflow(
+                v0_workflow_protos.WorkflowVersion,
+                self.v0_client.ScheduleWorkflow(
                     request,
                     metadata=get_metadata(self.token),
                 ),
@@ -286,6 +254,44 @@ class AdminClient:
 
             raise e
 
+    def _create_workflow_run_request(
+        self,
+        workflow_name: str,
+        input: JSONSerializableMapping,
+        options: TriggerWorkflowOptions,
+    ) -> v0_workflow_protos.TriggerWorkflowRequest:
+        workflow_run_id = ctx_workflow_run_id.get()
+        step_run_id = ctx_step_run_id.get()
+        worker_id = ctx_worker_id.get()
+        spawn_index = workflow_spawn_indices[workflow_run_id] if workflow_run_id else 0
+
+        ## Increment the spawn_index for the parent workflow
+        if workflow_run_id:
+            workflow_spawn_indices[workflow_run_id] += 1
+
+        desired_worker_id = (
+            (options.desired_worker_id or worker_id) if options.sticky else None
+        )
+        child_index = (
+            options.child_index if options.child_index is not None else spawn_index
+        )
+
+        trigger_options = TriggerWorkflowOptions(
+            parent_id=options.parent_id or workflow_run_id,
+            parent_step_run_id=options.parent_step_run_id or step_run_id,
+            child_key=options.child_key,
+            child_index=child_index,
+            additional_metadata=options.additional_metadata,
+            desired_worker_id=desired_worker_id,
+        )
+
+        namespace = options.namespace or self.namespace
+
+        if namespace != "" and not workflow_name.startswith(self.namespace):
+            workflow_name = f"{namespace}{workflow_name}"
+
+        return self._prepare_workflow_request(workflow_name, input, trigger_options)
+
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     @tenacity_retry
     def run_workflow(
@@ -294,70 +300,125 @@ class AdminClient:
         input: JSONSerializableMapping,
         options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
     ) -> WorkflowRunRef:
+        request = self._create_workflow_run_request(workflow_name, input, options)
+
+        if not self.pooled_workflow_listener:
+            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
         try:
-            if not self.pooled_workflow_listener:
-                self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+            resp = cast(
+                v0_workflow_protos.TriggerWorkflowResponse,
+                self.v0_client.TriggerWorkflow(
+                    request,
+                    metadata=get_metadata(self.token),
+                ),
+            )
+        except (grpc.RpcError, grpc.aio.AioRpcError) as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise DedupeViolationErr(e.details())
 
-            namespace = options.namespace or self.namespace
+        return WorkflowRunRef(
+            workflow_run_id=resp.workflow_run_id,
+            workflow_listener=self.pooled_workflow_listener,
+            workflow_run_event_listener=self.listener_client,
+        )
 
-            if namespace != "" and not workflow_name.startswith(self.namespace):
-                workflow_name = f"{namespace}{workflow_name}"
+    ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
+    @tenacity_retry
+    async def aio_run_workflow(
+        self,
+        workflow_name: str,
+        input: JSONSerializableMapping,
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> WorkflowRunRef:
+        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
+        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
+        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
+        async with spawn_index_lock:
+            request = self._create_workflow_run_request(workflow_name, input, options)
 
-            request = self._prepare_workflow_request(workflow_name, input, options)
+        if not self.pooled_workflow_listener:
+            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
 
+        try:
             resp = cast(
-                workflow_protos.TriggerWorkflowResponse,
-                self.client.TriggerWorkflow(
+                v0_workflow_protos.TriggerWorkflowResponse,
+                self.v0_client.TriggerWorkflow(
                     request,
                     metadata=get_metadata(self.token),
                 ),
             )
-
-            return WorkflowRunRef(
-                workflow_run_id=resp.workflow_run_id,
-                workflow_listener=self.pooled_workflow_listener,
-                workflow_run_event_listener=self.listener_client,
-            )
         except (grpc.RpcError, grpc.aio.AioRpcError) as e:
             if e.code() == grpc.StatusCode.ALREADY_EXISTS:
                 raise DedupeViolationErr(e.details())
 
             raise e
 
-    def _prepare_workflow_run_request(
-        self, workflow: WorkflowRunDict, options: TriggerWorkflowOptions
-    ) -> workflow_protos.TriggerWorkflowRequest:
-        workflow_name = workflow.workflow_name
-        input_data = workflow.input
-        options = workflow.options
-
-        namespace = options.namespace or self.namespace
-
-        if namespace != "" and not workflow_name.startswith(self.namespace):
-            workflow_name = f"{namespace}{workflow_name}"
-
-        return self._prepare_workflow_request(workflow_name, input_data, options)
+        return WorkflowRunRef(
+            workflow_run_id=resp.workflow_run_id,
+            workflow_listener=self.pooled_workflow_listener,
+            workflow_run_event_listener=self.listener_client,
+        )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     @tenacity_retry
     def run_workflows(
         self,
-        workflows: list[WorkflowRunDict],
-        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+        workflows: list[WorkflowRunTriggerConfig],
     ) -> list[WorkflowRunRef]:
         if not self.pooled_workflow_listener:
             self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
 
-        bulk_request = workflow_protos.BulkTriggerWorkflowRequest(
+        bulk_request = v0_workflow_protos.BulkTriggerWorkflowRequest(
             workflows=[
-                self._prepare_workflow_run_request(workflow, options)
+                self._create_workflow_run_request(
+                    workflow.workflow_name, workflow.input, workflow.options
+                )
                 for workflow in workflows
             ]
         )
 
         resp = cast(
-            workflow_protos.BulkTriggerWorkflowResponse,
-            self.client.BulkTriggerWorkflow(
+            v0_workflow_protos.BulkTriggerWorkflowResponse,
+            self.v0_client.BulkTriggerWorkflow(
+                bulk_request,
+                metadata=get_metadata(self.token),
+            ),
+        )
+
+        return [
+            WorkflowRunRef(
+                workflow_run_id=workflow_run_id,
+                workflow_listener=self.pooled_workflow_listener,
+                workflow_run_event_listener=self.listener_client,
+            )
+            for workflow_run_id in resp.workflow_run_ids
+        ]
+
+    @tenacity_retry
+    async def aio_run_workflows(
+        self,
+        workflows: list[WorkflowRunTriggerConfig],
+    ) -> list[WorkflowRunRef]:
+        ## IMPORTANT: The `pooled_workflow_listener` must be created 1) lazily, and not at `init` time, and 2) on the
+        ## main thread. If 1) is not followed, you'll get an error about something being attached to the wrong event
+        ## loop. If 2) is not followed, you'll get an error about the event loop not being set up.
+        if not self.pooled_workflow_listener:
+            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
+        async with spawn_index_lock:
+            bulk_request = v0_workflow_protos.BulkTriggerWorkflowRequest(
+                workflows=[
+                    self._create_workflow_run_request(
+                        workflow.workflow_name, workflow.input, workflow.options
+                    )
+                    for workflow in workflows
+                ]
+            )
+
+        resp = cast(
+            v0_workflow_protos.BulkTriggerWorkflowResponse,
+            self.v0_client.BulkTriggerWorkflow(
                 bulk_request,
                 metadata=get_metadata(self.token),
             ),

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -1,14 +1,14 @@
 import asyncio
 import json
 import time
-from dataclasses import dataclass, field
+from dataclasses import field
 from enum import Enum
-from typing import Any, AsyncGenerator, Optional, cast
+from typing import Any, AsyncGenerator, cast
 
 import grpc
 import grpc.aio
 from grpc._cython import cygrpc  # type: ignore[attr-defined]
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
 from hatchet_sdk.clients.events import proto_timestamp_now
@@ -36,25 +36,29 @@ DEFAULT_ACTION_TIMEOUT = 600  # seconds
 DEFAULT_ACTION_LISTENER_RETRY_COUNT = 15
 
 
-@dataclass
-class GetActionListenerRequest:
+class GetActionListenerRequest(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     worker_name: str
     services: list[str]
     actions: list[str]
-    max_runs: int | None = None
-    _labels: dict[str, str | int] = field(default_factory=dict)
+    slots: int = 100
+    raw_labels: dict[str, str | int] = Field(default_factory=dict)
 
-    labels: dict[str, WorkerLabels] = field(init=False)
+    labels: dict[str, WorkerLabels] = Field(default_factory=dict)
 
-    def __post_init__(self) -> None:
+    @model_validator(mode="after")
+    def validate_labels(self) -> "GetActionListenerRequest":
         self.labels = {}
 
-        for key, value in self._labels.items():
+        for key, value in self.raw_labels.items():
             if isinstance(value, int):
                 self.labels[key] = WorkerLabels(intValue=value)
             else:
                 self.labels[key] = WorkerLabels(strValue=str(value))
 
+        return self
+
 
 class ActionPayload(BaseModel):
     model_config = ConfigDict(extra="allow")
@@ -65,6 +69,7 @@ class ActionPayload(BaseModel):
     user_data: JSONSerializableMapping = Field(default_factory=dict)
     step_run_errors: dict[str, str] = Field(default_factory=dict)
     triggered_by: str | None = None
+    triggers: JSONSerializableMapping = Field(default_factory=dict)
 
     @field_validator(
         "input", "parents", "overrides", "user_data", "step_run_errors", mode="before"
@@ -142,30 +147,24 @@ def parse_additional_metadata(additional_metadata: str) -> JSONSerializableMappi
         return {}
 
 
-@dataclass
 class ActionListener:
-    config: ClientConfig
-    worker_id: str
+    def __init__(self, config: ClientConfig, worker_id: str) -> None:
+        self.config = config
+        self.worker_id = worker_id
 
-    client: DispatcherStub = field(init=False)
-    aio_client: DispatcherStub = field(init=False)
-    token: str = field(init=False)
-    retries: int = field(default=0, init=False)
-    last_connection_attempt: float = field(default=0, init=False)
-    last_heartbeat_succeeded: bool = field(default=True, init=False)
-    time_last_hb_succeeded: float = field(default=9999999999999, init=False)
-    heartbeat_task: Optional[asyncio.Task[None]] = field(default=None, init=False)
-    run_heartbeat: bool = field(default=True, init=False)
-    listen_strategy: str = field(default="v2", init=False)
-    stop_signal: bool = field(default=False, init=False)
-
-    missed_heartbeats: int = field(default=0, init=False)
-
-    def __post_init__(self) -> None:
-        self.client = DispatcherStub(new_conn(self.config, False))  # type: ignore[no-untyped-call]
         self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
         self.token = self.config.token
 
+        self.retries = 0
+        self.last_heartbeat_succeeded = True
+        self.time_last_hb_succeeded = 9999999999999.0
+        self.last_connection_attempt = 0.0
+        self.heartbeat_task: asyncio.Task[None] | None = None
+        self.run_heartbeat = True
+        self.listen_strategy = "v2"
+        self.stop_signal = False
+        self.missed_heartbeats = 0
+
     def is_healthy(self) -> bool:
         return self.last_heartbeat_succeeded
 
@@ -292,11 +291,16 @@ class ActionListener:
 
                     self.retries = 0
 
-                    action_payload = (
-                        {}
-                        if not assigned_action.actionPayload
-                        else self.parse_action_payload(assigned_action.actionPayload)
-                    )
+                    try:
+                        action_payload = (
+                            ActionPayload()
+                            if not assigned_action.actionPayload
+                            else ActionPayload.model_validate_json(
+                                assigned_action.actionPayload
+                            )
+                        )
+                    except (ValueError, json.JSONDecodeError) as e:
+                        raise ValueError(f"Error decoding payload: {e}")
 
                     action = Action(
                         tenant_id=assigned_action.tenantId,
@@ -309,7 +313,7 @@ class ActionListener:
                         step_id=assigned_action.stepId,
                         step_run_id=assigned_action.stepRunId,
                         action_id=assigned_action.actionId,
-                        action_payload=ActionPayload.model_validate(action_payload),
+                        action_payload=action_payload,
                         action_type=convert_proto_enum_to_python(
                             assigned_action.actionType,
                             ActionType,
@@ -352,16 +356,10 @@ class ActionListener:
 
                     self.retries = self.retries + 1
 
-    def parse_action_payload(self, payload: str) -> JSONSerializableMapping:
-        try:
-            return cast(JSONSerializableMapping, json.loads(payload))
-        except json.JSONDecodeError as e:
-            raise ValueError(f"Error decoding payload: {e}")
-
     async def get_listen_client(
         self,
     ) -> grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction]:
-        current_time = int(time.time())
+        current_time = time.time()
 
         if (
             current_time - self.last_connection_attempt
@@ -438,8 +436,10 @@ class ActionListener:
                 timeout=5,
                 metadata=get_metadata(self.token),
             )
+
             if self.interrupt is not None:
                 self.interrupt.set()
+
             return cast(WorkerUnsubscribeRequest, req)
         except grpc.RpcError as e:
             raise Exception(f"Failed to unsubscribe: {e}")