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

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -2,18 +2,23 @@ import asyncio
 import json
 import time
 from dataclasses import dataclass, field
-from typing import Any, AsyncGenerator, List, Optional
+from enum import Enum
+from typing import Any, AsyncGenerator, Optional, cast
 
 import grpc
-from grpc._cython import cygrpc
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
-from hatchet_sdk.clients.event_ts import Event_ts, read_with_interrupt
+from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
+from hatchet_sdk.clients.events import proto_timestamp_now
 from hatchet_sdk.clients.run_event_listener import (
     DEFAULT_ACTION_LISTENER_RETRY_INTERVAL,
 )
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
+from hatchet_sdk.contracts.dispatcher_pb2 import ActionType as ActionTypeProto
 from hatchet_sdk.contracts.dispatcher_pb2 import (
-    ActionType,
     AssignedAction,
     HeartbeatRequest,
     WorkerLabels,
@@ -22,31 +27,26 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
 )
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
 from hatchet_sdk.logger import logger
+from hatchet_sdk.metadata import get_metadata
 from hatchet_sdk.utils.backoff import exp_backoff_sleep
-from hatchet_sdk.utils.serialization import flatten
-
-from ...loader import ClientConfig
-from ...metadata import get_metadata
-from ..events import proto_timestamp_now
+from hatchet_sdk.utils.proto_enums import convert_proto_enum_to_python
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 DEFAULT_ACTION_TIMEOUT = 600  # seconds
-
-
-DEFAULT_ACTION_LISTENER_RETRY_INTERVAL = 5  # seconds
 DEFAULT_ACTION_LISTENER_RETRY_COUNT = 15
 
 
 @dataclass
 class GetActionListenerRequest:
     worker_name: str
-    services: List[str]
-    actions: List[str]
-    max_runs: Optional[int] = None
+    services: list[str]
+    actions: list[str]
+    max_runs: int | None = None
     _labels: dict[str, str | int] = field(default_factory=dict)
 
     labels: dict[str, WorkerLabels] = field(init=False)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self.labels = {}
 
         for key, value in self._labels.items():
@@ -56,8 +56,31 @@ class GetActionListenerRequest:
                 self.labels[key] = WorkerLabels(strValue=str(value))
 
 
-@dataclass
-class Action:
+class ActionPayload(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    input: JSONSerializableMapping = Field(default_factory=dict)
+    parents: dict[str, JSONSerializableMapping] = Field(default_factory=dict)
+    overrides: JSONSerializableMapping = Field(default_factory=dict)
+    user_data: JSONSerializableMapping = Field(default_factory=dict)
+    step_run_errors: dict[str, str] = Field(default_factory=dict)
+    triggered_by: str | None = None
+
+    @field_validator(
+        "input", "parents", "overrides", "user_data", "step_run_errors", mode="before"
+    )
+    @classmethod
+    def validate_fields(cls, v: Any) -> Any:
+        return v or {}
+
+
+class ActionType(str, Enum):
+    START_STEP_RUN = "START_STEP_RUN"
+    CANCEL_STEP_RUN = "CANCEL_STEP_RUN"
+    START_GET_GROUP_KEY = "START_GET_GROUP_KEY"
+
+
+class Action(BaseModel):
     worker_id: str
     tenant_id: str
     workflow_run_id: str
@@ -68,31 +91,25 @@ class Action:
     step_id: str
     step_run_id: str
     action_id: str
-    action_payload: str
     action_type: ActionType
     retry_count: int
-    additional_metadata: dict[str, str] | None = None
+    action_payload: ActionPayload
+    additional_metadata: JSONSerializableMapping = field(default_factory=dict)
 
     child_workflow_index: int | None = None
     child_workflow_key: str | None = None
     parent_workflow_run_id: str | None = None
 
-    def __post_init__(self):
-        if isinstance(self.additional_metadata, str) and self.additional_metadata != "":
-            try:
-                self.additional_metadata = json.loads(self.additional_metadata)
-            except json.JSONDecodeError:
-                # If JSON decoding fails, keep the original string
-                pass
-
-        # Ensure additional_metadata is always a dictionary
-        if not isinstance(self.additional_metadata, dict):
-            self.additional_metadata = {}
+    def _dump_payload_to_str(self) -> str:
+        try:
+            return json.dumps(self.action_payload.model_dump(), default=str)
+        except Exception:
+            return str(self.action_payload)
 
     @property
     def otel_attributes(self) -> dict[str, str | int]:
         try:
-            payload_str = json.dumps(self.action_payload, default=str)
+            payload_str = json.dumps(self.action_payload.model_dump(), default=str)
         except Exception:
             payload_str = str(self.action_payload)
 
@@ -115,9 +132,14 @@ class Action:
         return {k: v for k, v in attrs.items() if v}
 
 
-START_STEP_RUN = 0
-CANCEL_STEP_RUN = 1
-START_GET_GROUP_KEY = 2
+def parse_additional_metadata(additional_metadata: str) -> JSONSerializableMapping:
+    try:
+        return cast(
+            JSONSerializableMapping,
+            json.loads(additional_metadata),
+        )
+    except json.JSONDecodeError:
+        return {}
 
 
 @dataclass
@@ -132,22 +154,22 @@ class ActionListener:
     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] = field(default=None, 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):
-        self.client = DispatcherStub(new_conn(self.config))
-        self.aio_client = DispatcherStub(new_conn(self.config, True))
+    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
 
-    def is_healthy(self):
+    def is_healthy(self) -> bool:
         return self.last_heartbeat_succeeded
 
-    async def heartbeat(self):
+    async def heartbeat(self) -> None:
         # send a heartbeat every 4 seconds
         heartbeat_delay = 4
 
@@ -207,7 +229,7 @@ class ActionListener:
                     break
             await asyncio.sleep(heartbeat_delay)
 
-    async def start_heartbeater(self):
+    async def start_heartbeater(self) -> None:
         if self.heartbeat_task is not None:
             return
 
@@ -221,10 +243,10 @@ class ActionListener:
                 raise e
         self.heartbeat_task = loop.create_task(self.heartbeat())
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncGenerator[Action | None, None]:
         return self._generator()
 
-    async def _generator(self) -> AsyncGenerator[Action, None]:
+    async def _generator(self) -> AsyncGenerator[Action | None, None]:
         listener = None
 
         while not self.stop_signal:
@@ -239,7 +261,11 @@ class ActionListener:
 
             try:
                 while not self.stop_signal:
-                    self.interrupt = Event_ts()
+                    self.interrupt = ThreadSafeEvent()
+
+                    if listener is None:
+                        continue
+
                     t = asyncio.create_task(
                         read_with_interrupt(listener, self.interrupt)
                     )
@@ -252,7 +278,10 @@ class ActionListener:
                         )
 
                         t.cancel()
-                        listener.cancel()
+
+                        if listener:
+                            listener.cancel()
+
                         break
 
                     assigned_action = t.result()
@@ -262,20 +291,12 @@ class ActionListener:
                         break
 
                     self.retries = 0
-                    assigned_action: AssignedAction
-
-                    # Process the received action
-                    action_type = self.map_action_type(assigned_action.actionType)
 
-                    if (
-                        assigned_action.actionPayload is None
-                        or assigned_action.actionPayload == ""
-                    ):
-                        action_payload = None
-                    else:
-                        action_payload = self.parse_action_payload(
-                            assigned_action.actionPayload
-                        )
+                    action_payload = (
+                        {}
+                        if not assigned_action.actionPayload
+                        else self.parse_action_payload(assigned_action.actionPayload)
+                    )
 
                     action = Action(
                         tenant_id=assigned_action.tenantId,
@@ -288,10 +309,16 @@ class ActionListener:
                         step_id=assigned_action.stepId,
                         step_run_id=assigned_action.stepRunId,
                         action_id=assigned_action.actionId,
-                        action_payload=action_payload,
-                        action_type=action_type,
+                        action_payload=ActionPayload.model_validate(action_payload),
+                        action_type=convert_proto_enum_to_python(
+                            assigned_action.actionType,
+                            ActionType,
+                            ActionTypeProto,
+                        ),
                         retry_count=assigned_action.retryCount,
-                        additional_metadata=assigned_action.additional_metadata,
+                        additional_metadata=parse_additional_metadata(
+                            assigned_action.additional_metadata
+                        ),
                         child_workflow_index=assigned_action.child_workflow_index,
                         child_workflow_key=assigned_action.child_workflow_key,
                         parent_workflow_run_id=assigned_action.parent_workflow_run_id,
@@ -325,25 +352,15 @@ class ActionListener:
 
                     self.retries = self.retries + 1
 
-    def parse_action_payload(self, payload: str):
+    def parse_action_payload(self, payload: str) -> JSONSerializableMapping:
         try:
-            payload_data = json.loads(payload)
+            return cast(JSONSerializableMapping, json.loads(payload))
         except json.JSONDecodeError as e:
             raise ValueError(f"Error decoding payload: {e}")
-        return payload_data
-
-    def map_action_type(self, action_type):
-        if action_type == ActionType.START_STEP_RUN:
-            return START_STEP_RUN
-        elif action_type == ActionType.CANCEL_STEP_RUN:
-            return CANCEL_STEP_RUN
-        elif action_type == ActionType.START_GET_GROUP_KEY:
-            return START_GET_GROUP_KEY
-        else:
-            # logger.error(f"Unknown action type: {action_type}")
-            return None
 
-    async def get_listen_client(self):
+    async def get_listen_client(
+        self,
+    ) -> grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction]:
         current_time = int(time.time())
 
         if (
@@ -371,7 +388,7 @@ class ActionListener:
                 f"action listener connection interrupted, retrying... ({self.retries}/{DEFAULT_ACTION_LISTENER_RETRY_COUNT})"
             )
 
-        self.aio_client = DispatcherStub(new_conn(self.config, True))
+        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
 
         if self.listen_strategy == "v2":
             # we should await for the listener to be established before
@@ -392,11 +409,14 @@ class ActionListener:
 
         self.last_connection_attempt = current_time
 
-        return listener
+        return cast(
+            grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction], listener
+        )
 
-    def cleanup(self):
+    def cleanup(self) -> None:
         self.run_heartbeat = False
-        self.heartbeat_task.cancel()
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
 
         try:
             self.unregister()
@@ -406,9 +426,11 @@ class ActionListener:
         if self.interrupt:
             self.interrupt.set()
 
-    def unregister(self):
+    def unregister(self) -> WorkerUnsubscribeRequest:
         self.run_heartbeat = False
-        self.heartbeat_task.cancel()
+
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
 
         try:
             req = self.aio_client.Unsubscribe(
@@ -418,6 +440,6 @@ class ActionListener:
             )
             if self.interrupt is not None:
                 self.interrupt.set()
-            return req
+            return cast(WorkerUnsubscribeRequest, req)
         except grpc.RpcError as e:
             raise Exception(f"Failed to unsubscribe: {e}")

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -1,5 +1,6 @@
-from typing import Any, cast
+from typing import cast
 
+import grpc.aio
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from hatchet_sdk.clients.dispatcher.action_listener import (
@@ -8,6 +9,7 @@ from hatchet_sdk.clients.dispatcher.action_listener import (
     GetActionListenerRequest,
 )
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     STEP_EVENT_TYPE_COMPLETED,
@@ -26,22 +28,16 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
     WorkerRegisterResponse,
 )
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
-
-from ...loader import ClientConfig
-from ...metadata import get_metadata
+from hatchet_sdk.metadata import get_metadata
 
 DEFAULT_REGISTER_TIMEOUT = 30
 
 
-def new_dispatcher(config: ClientConfig) -> "DispatcherClient":
-    return DispatcherClient(config=config)
-
-
 class DispatcherClient:
     config: ClientConfig
 
     def __init__(self, config: ClientConfig):
-        conn = new_conn(config)
+        conn = new_conn(config, False)
         self.client = DispatcherStub(conn)  # type: ignore[no-untyped-call]
 
         aio_conn = new_conn(config, True)
@@ -76,7 +72,7 @@ class DispatcherClient:
 
     async def send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse] | None:
         try:
             return await self._try_send_step_action_event(action, event_type, payload)
         except Exception as e:
@@ -91,12 +87,12 @@ class DispatcherClient:
                     "Failed to send finished event: " + str(e),
                 )
 
-            return
+            return None
 
     @tenacity_retry
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse]:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -113,15 +109,17 @@ class DispatcherClient:
             retryCount=action.retry_count,
         )
 
-        ## TODO: What does this return?
-        return await self.aio_client.SendStepActionEvent(
-            event,
-            metadata=get_metadata(self.token),
+        return cast(
+            grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse],
+            await self.aio_client.SendStepActionEvent(
+                event,
+                metadata=get_metadata(self.token),
+            ),
         )
 
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse]:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -135,10 +133,12 @@ class DispatcherClient:
             eventPayload=payload,
         )
 
-        ## TODO: What does this return?
-        return await self.aio_client.SendGroupKeyActionEvent(
-            event,
-            metadata=get_metadata(self.token),
+        return cast(
+            grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse],
+            await self.aio_client.SendGroupKeyActionEvent(
+                event,
+                metadata=get_metadata(self.token),
+            ),
         )
 
     def put_overrides_data(self, data: OverridesData) -> ActionEventResponse:

hatchet_sdk/clients/event_ts.py

@@ -1,28 +1,41 @@
 import asyncio
-from typing import Any
+from typing import TypeVar, cast
 
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
-class Event_ts(asyncio.Event):
+
+class ThreadSafeEvent(asyncio.Event):
     """
-    Event_ts is a subclass of asyncio.Event that allows for thread-safe setting and clearing of the event.
+    ThreadSafeEvent is a subclass of asyncio.Event that allows for thread-safe setting and clearing of the event.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        if self._loop is None:
+    def __init__(self) -> None:
+        super().__init__()
+        if self._loop is None:  # type: ignore[has-type]
             self._loop = asyncio.get_event_loop()
 
-    def set(self):
+    def set(self) -> None:
         if not self._loop.is_closed():
             self._loop.call_soon_threadsafe(super().set)
 
-    def clear(self):
+    def clear(self) -> None:
         self._loop.call_soon_threadsafe(super().clear)
 
 
-async def read_with_interrupt(listener: Any, interrupt: Event_ts):
+TRequest = TypeVar("TRequest")
+TResponse = TypeVar("TResponse")
+
+
+async def read_with_interrupt(
+    listener: grpc.aio.UnaryStreamCall[TRequest, TResponse], interrupt: ThreadSafeEvent
+) -> TResponse:
     try:
         result = await listener.read()
-        return result
+
+        if result is cygrpc.EOF:
+            raise ValueError("Unexpected EOF")
+
+        return cast(TResponse, result)
     finally:
         interrupt.set()