more cleanup

Author: adrianlyjak

src/workflows/context/context.py

@@ -153,7 +153,7 @@ def __init__(
             if (
                 step_config.context_state_type is not None
                 and step_config.context_state_type != DictState
-                and isinstance(step_config.context_state_type, type(BaseModel))
+                and issubclass(step_config.context_state_type, BaseModel)
             ):
                 state_type = step_config.context_state_type
                 state_types.add(state_type)
@@ -207,14 +207,14 @@ def _init_broker(self, workflow: Workflow) -> WorkflowBroker[MODEL_T]:
         )
         return self._broker_run
 
-    def _internal_run(
+    def _workflow_run(
         self,
         workflow: Workflow,
         start_event: StartEvent | None = None,
         semaphore: asyncio.Semaphore | None = None,
     ) -> WorkflowHandler:
         """
-        Called internally from the workflow to run it
+        called by package internally from the workflow to run it
         """
         prev_broker: WorkflowBroker[MODEL_T] | None = None
         if self._broker_run is not None:
@@ -243,7 +243,7 @@ async def after_complete() -> None:
             after_complete=after_complete,
         )
 
-    def _internal_cancel_run(self) -> None:
+    def _workflow_cancel_run(self) -> None:
         """
         Called internally from the handler to cancel a context's run
         """

src/workflows/decorators.py

@@ -55,7 +55,7 @@ class StepConfig(BaseModel):
 
 
 class StepFunction(Protocol, Generic[P, R]):
-    """A decorated function, that has some step_config metadata from the @step decorator"""
+    """A decorated function, that has some _step_config metadata from the @step decorator"""
 
     _step_config: StepConfig
 

src/workflows/handler.py

@@ -6,7 +6,7 @@
 import asyncio
 from typing import Any, AsyncGenerator
 
-from .context.context import Context
+from .context import Context
 from .errors import WorkflowRuntimeError
 from .events import Event, StopEvent, InternalDispatchEvent
 from .types import RunResultT
@@ -129,6 +129,9 @@ async def cancel_run(self) -> None:
             ```
         """
         if self.ctx:
-            self.ctx._internal_cancel_run()
+            self.ctx._workflow_cancel_run()
             if self._run_task is not None:
-                await self._run_task
+                try:
+                    await self._run_task
+                except Exception:
+                    pass

src/workflows/runtime/broker.py

@@ -124,8 +124,14 @@ def __init__(
         # Store the step configs of this workflow, to be used in send_event
         self._step_configs = {}
 
+        # Make the system step "_done" accept custom stop events
+        done_step = self._done
+        done_accepted_events = done_step._step_config.accepted_events
+        if workflow._stop_event_class not in done_accepted_events:
+            done_accepted_events.append(workflow._stop_event_class)
+
         steps: dict[str, StepFunction] = {
-            "_done": self._done,
+            "_done": done_step,
             **workflow._get_steps(),
         }
         for step_name, step_func in steps.items():
@@ -161,13 +167,6 @@ def __init__(
 
             step_config: StepConfig = step_func._step_config
 
-            # Make the system step "_done" accept custom stop events
-            if (
-                name == "_done"
-                and workflow._stop_event_class not in step_config.accepted_events
-            ):
-                step_config.accepted_events.append(workflow._stop_event_class)
-
             for _ in range(step_config.num_workers):
                 self._add_step_worker(
                     name=name,
@@ -220,6 +219,9 @@ def start(
             )
 
         async def _run_workflow() -> None:
+            # defer execution to make sure the task can be captured and passed
+            # to the handler, protecting against exceptions from before_start
+            await asyncio.sleep(0)
             if before_start is not None:
                 await before_start()
             try:

src/workflows/testing/runner.py

@@ -67,15 +67,11 @@ async def run(
         """
         handler = self._workflow.run(start_event=start_event, ctx=ctx)
         collected_events: list[Event] = []
-        print("streaming events")
         async for event in handler.stream_events(expose_internal=expose_internal):
-            print("streaming event", type(event).__name__)
             if exclude_events and type(event) in exclude_events:
                 continue
             collected_events.append(event)
-        print("awaiting handler")
         result = await handler
-        print("result", result)
         event_freqs: dict[EventType, int] = dict(
             Counter([type(ev) for ev in collected_events])
         )

src/workflows/workflow.py

@@ -173,7 +173,7 @@ def _ensure_events_collected(self) -> list[type[Event]]:
         """
         events_found: set[type[Event]] = set()
         for step_func in self._get_steps().values():
-            step_config: StepConfig = getattr(step_func, "__step_config")
+            step_config: StepConfig = step_func._step_config
 
             # Do not collect events from the done step
             if step_func.__name__ == "_done":
@@ -332,7 +332,7 @@ def run(
             if ctx.is_running
             else self._get_start_event_instance(start_event, **kwargs)
         )
-        return ctx._internal_run(
+        return ctx._workflow_run(
             workflow=self, start_event=start_event_instance, semaphore=self._sem
         )