hatchet-sdk 0.47.1__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

hatchet_sdk/clients/durable_event_listener.py

@@ -0,0 +1,327 @@
+import asyncio
+import json
+from collections.abc import AsyncIterator
+from typing import Any, Literal, cast
+
+import grpc
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
+from pydantic import BaseModel, ConfigDict
+
+from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
+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.v1.dispatcher_pb2 import (
+    DurableEvent,
+    ListenForDurableEventRequest,
+)
+from hatchet_sdk.contracts.v1.dispatcher_pb2 import (
+    RegisterDurableEventRequest as RegisterDurableEventRequestProto,
+)
+from hatchet_sdk.contracts.v1.dispatcher_pb2_grpc import V1DispatcherStub
+from hatchet_sdk.contracts.v1.shared.condition_pb2 import DurableEventListenerConditions
+from hatchet_sdk.logger import logger
+from hatchet_sdk.metadata import get_metadata
+from hatchet_sdk.waits import SleepCondition, UserEventCondition
+
+DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL = 3  # seconds
+DEFAULT_DURABLE_EVENT_LISTENER_RETRY_COUNT = 5
+DEFAULT_DURABLE_EVENT_LISTENER_INTERRUPT_INTERVAL = 1800  # 30 minutes
+
+
+class _Subscription:
+    def __init__(self, id: int, task_id: str, signal_key: str):
+        self.id = id
+        self.task_id = task_id
+        self.signal_key = signal_key
+        self.queue: asyncio.Queue[DurableEvent | None] = asyncio.Queue()
+
+    async def __aiter__(self) -> "_Subscription":
+        return self
+
+    async def __anext__(self) -> DurableEvent | None:
+        return await self.queue.get()
+
+    async def get(self) -> DurableEvent:
+        event = await self.queue.get()
+
+        if event is None:
+            raise StopAsyncIteration
+
+        return event
+
+    async def put(self, item: DurableEvent) -> None:
+        await self.queue.put(item)
+
+    async def close(self) -> None:
+        await self.queue.put(None)
+
+
+class RegisterDurableEventRequest(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    task_id: str
+    signal_key: str
+    conditions: list[SleepCondition | UserEventCondition]
+
+    def to_proto(self) -> RegisterDurableEventRequestProto:
+        return RegisterDurableEventRequestProto(
+            task_id=self.task_id,
+            signal_key=self.signal_key,
+            conditions=DurableEventListenerConditions(
+                sleep_conditions=[
+                    c.to_pb() for c in self.conditions if isinstance(c, SleepCondition)
+                ],
+                user_event_conditions=[
+                    c.to_pb()
+                    for c in self.conditions
+                    if isinstance(c, UserEventCondition)
+                ],
+            ),
+        )
+
+
+class DurableEventListener:
+    def __init__(self, config: ClientConfig):
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        conn = new_conn(config, True)
+        self.client = V1DispatcherStub(conn)  # type: ignore[no-untyped-call]
+        self.token = config.token
+        self.config = config
+
+        # list of all active subscriptions, mapping from a subscription id to a task id and signal key
+        self.subscriptions_to_task_id_signal_key: dict[int, tuple[str, str]] = {}
+
+        # task id-signal key tuples mapped to an array of subscription ids
+        self.task_id_signal_key_to_subscriptions: dict[tuple[str, str], list[int]] = {}
+
+        self.subscription_counter: int = 0
+        self.subscription_counter_lock: asyncio.Lock = asyncio.Lock()
+
+        self.requests: asyncio.Queue[ListenForDurableEventRequest | int] = (
+            asyncio.Queue()
+        )
+
+        self.listener: (
+            grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent] | None
+        ) = None
+        self.listener_task: asyncio.Task[None] | None = None
+
+        self.curr_requester: int = 0
+
+        self.events: dict[int, _Subscription] = {}
+
+        self.interrupter: asyncio.Task[None] | None = None
+
+    async def _interrupter(self) -> None:
+        """
+        _interrupter runs in a separate thread and interrupts the listener according to a configurable duration.
+        """
+        await asyncio.sleep(DEFAULT_DURABLE_EVENT_LISTENER_INTERRUPT_INTERVAL)
+
+        if self.interrupt is not None:
+            self.interrupt.set()
+
+    async def _init_producer(self) -> None:
+        try:
+            if not self.listener:
+                while True:
+                    try:
+                        self.listener = await self._retry_subscribe()
+
+                        logger.debug("Workflow run listener connected.")
+
+                        # spawn an interrupter task
+                        if self.interrupter is not None and not self.interrupter.done():
+                            self.interrupter.cancel()
+
+                        self.interrupter = asyncio.create_task(self._interrupter())
+
+                        while True:
+                            self.interrupt = ThreadSafeEvent()
+                            if self.listener is None:
+                                continue
+
+                            t = asyncio.create_task(
+                                read_with_interrupt(self.listener, self.interrupt)
+                            )
+                            await self.interrupt.wait()
+
+                            if not t.done():
+                                logger.warning(
+                                    "Interrupted read_with_interrupt task of durable event listener"
+                                )
+
+                                t.cancel()
+                                if self.listener:
+                                    self.listener.cancel()
+                                await asyncio.sleep(
+                                    DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL
+                                )
+                                break
+
+                            event = t.result()
+
+                            if event is cygrpc.EOF:
+                                break
+
+                            # get a list of subscriptions for this task-signal pair
+                            subscriptions = (
+                                self.task_id_signal_key_to_subscriptions.get(
+                                    (event.task_id, event.signal_key), []
+                                )
+                            )
+
+                            for subscription_id in subscriptions:
+                                await self.events[subscription_id].put(event)
+
+                    except grpc.RpcError as e:
+                        logger.debug(f"grpc error in durable event listener: {e}")
+                        await asyncio.sleep(
+                            DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL
+                        )
+                        continue
+
+        except Exception as e:
+            logger.error(f"Error in durable event listener: {e}")
+
+            self.listener = None
+
+            # close all subscriptions
+            for subscription_id in self.events:
+                await self.events[subscription_id].close()
+
+            raise e
+
+    async def _request(self) -> AsyncIterator[ListenForDurableEventRequest]:
+        self.curr_requester = self.curr_requester + 1
+
+        # replay all existing subscriptions
+        for task_id, signal_key in set(
+            self.subscriptions_to_task_id_signal_key.values()
+        ):
+            yield ListenForDurableEventRequest(
+                task_id=task_id,
+                signal_key=signal_key,
+            )
+
+        while True:
+            request = await self.requests.get()
+
+            # if the request is an int which matches the current requester, then we should stop
+            if request == self.curr_requester:
+                break
+
+            # if we've gotten an int that doesn't match the current requester, then we should ignore it
+            if isinstance(request, int):
+                continue
+
+            yield request
+            self.requests.task_done()
+
+    def cleanup_subscription(self, subscription_id: int) -> None:
+        task_id_signal_key = self.subscriptions_to_task_id_signal_key[subscription_id]
+
+        if task_id_signal_key in self.task_id_signal_key_to_subscriptions:
+            self.task_id_signal_key_to_subscriptions[task_id_signal_key].remove(
+                subscription_id
+            )
+
+        del self.subscriptions_to_task_id_signal_key[subscription_id]
+        del self.events[subscription_id]
+
+    async def subscribe(self, task_id: str, signal_key: str) -> DurableEvent:
+        try:
+            # create a new subscription id, place a mutex on the counter
+            async with self.subscription_counter_lock:
+                self.subscription_counter += 1
+                subscription_id = self.subscription_counter
+
+            self.subscriptions_to_task_id_signal_key[subscription_id] = (
+                task_id,
+                signal_key,
+            )
+
+            if (task_id, signal_key) not in self.task_id_signal_key_to_subscriptions:
+                self.task_id_signal_key_to_subscriptions[(task_id, signal_key)] = [
+                    subscription_id
+                ]
+            else:
+                self.task_id_signal_key_to_subscriptions[(task_id, signal_key)].append(
+                    subscription_id
+                )
+
+            self.events[subscription_id] = _Subscription(
+                subscription_id, task_id, signal_key
+            )
+
+            await self.requests.put(
+                ListenForDurableEventRequest(
+                    task_id=task_id,
+                    signal_key=signal_key,
+                )
+            )
+
+            if not self.listener_task or self.listener_task.done():
+                self.listener_task = asyncio.create_task(self._init_producer())
+
+            return await self.events[subscription_id].get()
+        except asyncio.CancelledError:
+            raise
+        finally:
+            self.cleanup_subscription(subscription_id)
+
+    async def _retry_subscribe(
+        self,
+    ) -> grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent]:
+        retries = 0
+
+        while retries < DEFAULT_DURABLE_EVENT_LISTENER_RETRY_COUNT:
+            try:
+                if retries > 0:
+                    await asyncio.sleep(DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL)
+
+                # signal previous async iterator to stop
+                if self.curr_requester != 0:
+                    self.requests.put_nowait(self.curr_requester)
+
+                return cast(
+                    grpc.aio.UnaryStreamCall[
+                        ListenForDurableEventRequest, DurableEvent
+                    ],
+                    self.client.ListenForDurableEvent(
+                        self._request(),
+                        metadata=get_metadata(self.token),
+                    ),
+                )
+            except grpc.RpcError as e:
+                if e.code() == grpc.StatusCode.UNAVAILABLE:
+                    retries = retries + 1
+                else:
+                    raise ValueError(f"gRPC error: {e}")
+
+        raise ValueError("Failed to connect to durable event listener")
+
+    @tenacity_retry
+    def register_durable_event(
+        self, request: RegisterDurableEventRequest
+    ) -> Literal[True]:
+        self.client.RegisterDurableEvent(
+            request.to_proto(),
+            timeout=5,
+            metadata=get_metadata(self.token),
+        )
+
+        return True
+
+    @tenacity_retry
+    async def result(self, task_id: str, signal_key: str) -> dict[str, Any]:
+        event = await self.subscribe(task_id, signal_key)
+
+        return cast(dict[str, Any], json.loads(event.data.decode("utf-8")))

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()

hatchet_sdk/clients/events.py

@@ -1,13 +1,14 @@
 import asyncio
 import datetime
 import json
-from typing import Any, Dict, List, Optional, TypedDict
-from uuid import uuid4
+from typing import List, cast
 
 import grpc
 from google.protobuf import timestamp_pb2
+from pydantic import BaseModel, Field
 
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.events_pb2 import (
     BulkPushEventRequest,
     Event,
@@ -16,20 +17,18 @@ from hatchet_sdk.contracts.events_pb2 import (
     PutStreamEventRequest,
 )
 from hatchet_sdk.contracts.events_pb2_grpc import EventsServiceStub
-from hatchet_sdk.utils.serialization import flatten
+from hatchet_sdk.metadata import get_metadata
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 
-from ..loader import ClientConfig
-from ..metadata import get_metadata
 
-
-def new_event(conn, config: ClientConfig):
+def new_event(conn: grpc.Channel, config: ClientConfig) -> "EventClient":
     return EventClient(
-        client=EventsServiceStub(conn),
+        client=EventsServiceStub(conn),  # type: ignore[no-untyped-call]
         config=config,
     )
 
 
-def proto_timestamp_now():
+def proto_timestamp_now() -> timestamp_pb2.Timestamp:
     t = datetime.datetime.now().timestamp()
     seconds = int(t)
     nanos = int(t % 1 * 1e9)
@@ -37,19 +36,19 @@ def proto_timestamp_now():
     return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
 
 
-class PushEventOptions(TypedDict, total=False):
-    additional_metadata: Dict[str, str] | None = None
+class PushEventOptions(BaseModel):
+    additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
     namespace: str | None = None
 
 
-class BulkPushEventOptions(TypedDict, total=False):
+class BulkPushEventOptions(BaseModel):
     namespace: str | None = None
 
 
-class BulkPushEventWithMetadata(TypedDict, total=False):
+class BulkPushEventWithMetadata(BaseModel):
     key: str
-    payload: Any
-    additional_metadata: Optional[Dict[str, Any]]  # Optional metadata
+    payload: JSONSerializableMapping = Field(default_factory=dict)
+    additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
 
 
 class EventClient:
@@ -58,126 +57,124 @@ class EventClient:
         self.token = config.token
         self.namespace = config.namespace
 
-    async def async_push(
-        self, event_key, payload, options: Optional[PushEventOptions] = None
+    async def aio_push(
+        self,
+        event_key: str,
+        payload: JSONSerializableMapping,
+        options: PushEventOptions = PushEventOptions(),
     ) -> Event:
         return await asyncio.to_thread(
             self.push, event_key=event_key, payload=payload, options=options
         )
 
-    async def async_bulk_push(
+    async def aio_bulk_push(
         self,
-        events: List[BulkPushEventWithMetadata],
-        options: Optional[BulkPushEventOptions] = None,
+        events: list[BulkPushEventWithMetadata],
+        options: BulkPushEventOptions = BulkPushEventOptions(),
     ) -> List[Event]:
         return await asyncio.to_thread(self.bulk_push, events=events, options=options)
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     @tenacity_retry
-    def push(self, event_key, payload, options: PushEventOptions = None) -> Event:
-        namespace = self.namespace
-
-        if (
-            options is not None
-            and "namespace" in options
-            and options["namespace"] is not None
-        ):
-            namespace = options.pop("namespace")
-
+    def push(
+        self,
+        event_key: str,
+        payload: JSONSerializableMapping,
+        options: PushEventOptions = PushEventOptions(),
+    ) -> Event:
+        namespace = options.namespace or self.namespace
         namespaced_event_key = namespace + event_key
 
         try:
-            meta = dict() if options is None else options["additional_metadata"]
-            meta_bytes = None if meta is None else json.dumps(meta).encode("utf-8")
+            meta = options.additional_metadata
+            meta_bytes = None if meta is None else json.dumps(meta)
         except Exception as e:
             raise ValueError(f"Error encoding meta: {e}")
 
         try:
-            payload_bytes = json.dumps(payload).encode("utf-8")
-        except json.UnicodeEncodeError as e:
+            payload_str = json.dumps(payload)
+        except (TypeError, ValueError) as e:
             raise ValueError(f"Error encoding payload: {e}")
 
         request = PushEventRequest(
             key=namespaced_event_key,
-            payload=payload_bytes,
+            payload=payload_str,
             eventTimestamp=proto_timestamp_now(),
             additionalMetadata=meta_bytes,
         )
 
-        return self.client.Push(request, metadata=get_metadata(self.token))
+        return cast(Event, self.client.Push(request, metadata=get_metadata(self.token)))
+
+    def _create_push_event_request(
+        self,
+        event: BulkPushEventWithMetadata,
+        namespace: str,
+    ) -> PushEventRequest:
+        event_key = namespace + event.key
+        payload = event.payload
+
+        meta = event.additional_metadata
+
+        try:
+            meta_str = json.dumps(meta)
+        except Exception as e:
+            raise ValueError(f"Error encoding meta: {e}")
+
+        try:
+            serialized_payload = json.dumps(payload)
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"Error serializing payload: {e}")
+
+        return PushEventRequest(
+            key=event_key,
+            payload=serialized_payload,
+            eventTimestamp=proto_timestamp_now(),
+            additionalMetadata=meta_str,
+        )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     @tenacity_retry
     def bulk_push(
         self,
         events: List[BulkPushEventWithMetadata],
-        options: BulkPushEventOptions = None,
+        options: BulkPushEventOptions = BulkPushEventOptions(),
     ) -> List[Event]:
-        namespace = self.namespace
-
-        if (
-            options is not None
-            and "namespace" in options
-            and options["namespace"] is not None
-        ):
-            namespace = options.pop("namespace")
-
-        bulk_events = []
-        for event in events:
-            event_key = namespace + event["key"]
-            payload = event["payload"]
-
-            try:
-                meta = event.get("additional_metadata", {})
-                meta_bytes = json.dumps(meta).encode("utf-8") if meta else None
-            except Exception as e:
-                raise ValueError(f"Error encoding meta: {e}")
-
-            try:
-                payload_bytes = json.dumps(payload).encode("utf-8")
-            except json.UnicodeEncodeError as e:
-                raise ValueError(f"Error encoding payload: {e}")
-
-            request = PushEventRequest(
-                key=event_key,
-                payload=payload_bytes,
-                eventTimestamp=proto_timestamp_now(),
-                additionalMetadata=meta_bytes,
-            )
-            bulk_events.append(request)
-
-        bulk_request = BulkPushEventRequest(events=bulk_events)
+        namespace = options.namespace or self.namespace
+
+        bulk_request = BulkPushEventRequest(
+            events=[
+                self._create_push_event_request(event, namespace) for event in events
+            ]
+        )
 
         response = self.client.BulkPush(bulk_request, metadata=get_metadata(self.token))
 
-        return response.events
+        return cast(
+            list[Event],
+            response.events,
+        )
 
-    def log(self, message: str, step_run_id: str):
-        try:
-            request = PutLogRequest(
-                stepRunId=step_run_id,
-                createdAt=proto_timestamp_now(),
-                message=message,
-            )
+    def log(self, message: str, step_run_id: str) -> None:
+        request = PutLogRequest(
+            stepRunId=step_run_id,
+            createdAt=proto_timestamp_now(),
+            message=message,
+        )
 
-            self.client.PutLog(request, metadata=get_metadata(self.token))
-        except Exception as e:
-            raise ValueError(f"Error logging: {e}")
+        self.client.PutLog(request, metadata=get_metadata(self.token))
 
-    def stream(self, data: str | bytes, step_run_id: str):
-        try:
-            if isinstance(data, str):
-                data_bytes = data.encode("utf-8")
-            elif isinstance(data, bytes):
-                data_bytes = data
-            else:
-                raise ValueError("Invalid data type. Expected str, bytes, or file.")
-
-            request = PutStreamEventRequest(
-                stepRunId=step_run_id,
-                createdAt=proto_timestamp_now(),
-                message=data_bytes,
-            )
-            self.client.PutStreamEvent(request, metadata=get_metadata(self.token))
-        except Exception as e:
-            raise ValueError(f"Error putting stream event: {e}")
+    def stream(self, data: str | bytes, step_run_id: str) -> None:
+        if isinstance(data, str):
+            data_bytes = data.encode("utf-8")
+        elif isinstance(data, bytes):
+            data_bytes = data
+        else:
+            raise ValueError("Invalid data type. Expected str, bytes, or file.")
+
+        request = PutStreamEventRequest(
+            stepRunId=step_run_id,
+            createdAt=proto_timestamp_now(),
+            message=data_bytes,
+        )
+
+        self.client.PutStreamEvent(request, metadata=get_metadata(self.token))