Adding registry workaround

Author: adrianlyjak

src/workflows/plugins/dbos.py

@@ -16,10 +16,9 @@
     RegisteredWorkflow,
 )
 from workflows.runtime.types.internal_state import BrokerState
-from workflows.context.context import Context
 from workflows.runtime.types.step_function import StepWorkerFunction
 from workflows.runtime.types.ticks import WorkflowTick
-from workflows.decorators import R
+
 from workflows.workflow import Workflow
 
 
@@ -33,50 +32,31 @@ def register(
         self,
         workflow: Workflow,
         workflow_function: ControlLoopFunction,
-        steps: dict[str, StepWorkerFunction[R]],
+        steps: dict[str, StepWorkerFunction],
     ) -> RegisteredWorkflow | None:
         """
         Wrap the workflow control loop in a DBOS workflow so ticks are received via DBOS.recv
         and sent via DBOS.send, enabling durable orchestration.
         """
 
-
-
-        # DBOS Python supports async workflow functions; we wrap and return replacement.
-        # Note: We do not wrap individual step workers here; the control loop retains
-        # responsibility for step execution. This can be extended later to split steps
-        # into child workflows if needed.
-
         @DBOS.workflow()
         async def _dbos_control_loop(
             start_event: Event | None,
             init_state: BrokerState | None,
-            plugin: WorkflowRuntime,
-            context: Context,
-            step_workers: dict[str, StepWorkerFunction],
+            run_id: str,
         ) -> StopEvent:
-            # Ensure our runtime knows the workflow id for routing incoming ticks
-            assert isinstance(plugin, DBOSWorkflowRuntime)
-            # Pin a stable workflow id for this run using the runtime's run_id
-            with SetWorkflowID(plugin.run_id):
-                # Delegate to the original control loop function
-                return await workflow_function(  # type: ignore[misc]
-                    start_event, init_state, plugin, context, step_workers
-                )
+            with SetWorkflowID(run_id):
+                return await workflow_function(start_event, init_state, run_id)
 
         async def wrapper(
             start_event: Event | None,
             init_state: BrokerState | None,
-            plugin: WorkflowRuntime,
-            context: Context,
-            step_workers: dict[str, StepWorkerFunction],
+            run_id: str,
         ) -> StopEvent:
             # Call the DBOS workflow directly; DBOS will orchestrate execution
-            return await _dbos_control_loop(
-                start_event, init_state, plugin, context, step_workers
-            )
+            return await _dbos_control_loop(start_event, init_state, run_id)
 
-        return RegisteredWorkflow(workflow_function=wrapper, steps=steps)
+        return RegisteredWorkflow(workflow_function=_dbos_control_loop, steps=steps)
 
     def new_runtime(self, run_id: str) -> WorkflowRuntime:
         runtime: WorkflowRuntime = DBOSWorkflowRuntime(run_id)

src/workflows/runtime/broker.py

@@ -27,7 +27,7 @@
     StartEvent,
     StopEvent,
 )
-from workflows.runtime.control_loop import create_control_loop, rebuild_state_from_ticks
+from workflows.runtime.control_loop import control_loop, rebuild_state_from_ticks
 from workflows.runtime.types.internal_state import BrokerState
 from workflows.runtime.types.plugin import Plugin, WorkflowRuntime, as_snapshottable
 from workflows.runtime.types.results import (
@@ -39,7 +39,10 @@
     StepWorkerStateContextVar,
     WaitingForEvent,
 )
-from workflows.runtime.types.step_function import as_step_worker_function
+from workflows.runtime.types.step_function import (
+    StepWorkerFunction,
+    as_step_worker_function,
+)
 from workflows.runtime.types.ticks import TickAddEvent, TickCancelRun, WorkflowTick
 from workflows.runtime.workflow_registry import workflow_registry
 
@@ -118,8 +121,6 @@ def start(
         self._init_state = previous
 
         async def _run_workflow() -> None:
-            from workflows.context.context import Context
-
             # defer execution to make sure the task can be captured and passed
             # to the handler as async exception, protecting against exceptions from before_start
             self._is_running = True
@@ -132,27 +133,35 @@ async def _run_workflow() -> None:
                 try:
                     exception_raised = None
 
-                    step_workers = {}
+                    step_workers: dict[str, StepWorkerFunction] = {}
                     for name, step_func in workflow._get_steps().items():
                         # Avoid capturing a bound method (which retains the instance).
                         # If it's a bound method, extract the unbound function from the class.
                         unbound = getattr(step_func, "__func__", step_func)
                         step_workers[name] = as_step_worker_function(unbound)
 
-                    control_loop_fn = create_control_loop(
-                        workflow,
-                    )
                     registered = workflow_registry.get_registered_workflow(
-                        workflow, self._plugin, control_loop_fn, step_workers
+                        workflow, self._plugin, control_loop, step_workers
                     )
 
-                    workflow_result = await registered.workflow_function(
-                        start_event,
-                        init_state,
-                        self._runtime,
-                        cast(Context, self._context),
-                        registered.steps,
+                    # Register run context prior to invoking control loop
+                    workflow_registry.register_run(
+                        run_id=run_id,
+                        workflow=workflow,
+                        plugin=self._runtime,
+                        context=self._context,  # type: ignore
+                        steps=registered.steps,
                     )
+
+                    try:
+                        workflow_result = await registered.workflow_function(
+                            start_event,
+                            init_state,
+                            run_id,
+                        )
+                    finally:
+                        # ensure run context is cleaned up even on failure
+                        workflow_registry.delete_run(run_id)
                     result.set_result(
                         workflow_result.result
                         if type(workflow_result) is StopEvent

src/workflows/runtime/control_loop.py

@@ -8,7 +8,7 @@
 from copy import deepcopy
 import time
 from typing import TYPE_CHECKING
-import weakref
+
 
 from workflows.decorators import R
 from workflows.errors import (
@@ -41,7 +41,6 @@
     InternalStepWorkerState,
 )
 from workflows.runtime.types.plugin import (
-    ControlLoopFunction,
     WorkflowRuntime,
     as_snapshottable,
 )
@@ -69,6 +68,7 @@
 import logging
 
 from workflows.workflow import Workflow
+from workflows.runtime.workflow_registry import workflow_registry
 
 if TYPE_CHECKING:
     from workflows.context.context import Context
@@ -282,37 +282,23 @@ async def _pull() -> None:
             await self.cleanup_tasks()
 
 
-def create_control_loop(
-    workflow: Workflow,
-) -> ControlLoopFunction:
+async def control_loop(
+    start_event: Event | None,
+    init_state: BrokerState | None,
+    run_id: str,
+) -> StopEvent:
     """
-    Creates a control loop for a workflow run. Dependencies are provided as initial args, and returns a simple function
-    with only start event and an optional start state as an arg, that can be easily decorated.
-
-    Returns a function that can be called to start a workflow run.
+    The main async control loop for a workflow run.
     """
-
-    workflow_ref = weakref.ref(workflow)
-
-    async def control_loop(
-        start_event: Event | None,
-        init_state: BrokerState | None,
-        # TODO - get these 3 out of here! Needs to be inferred from scope somehow for proper distributed, static execution
-        plugin: WorkflowRuntime,
-        context: Context,
-        step_workers: dict[str, StepWorkerFunction],
-    ) -> StopEvent:
-        """
-        The main async control loop for a workflow run.
-        """
-        wf = workflow_ref()
-        if wf is None:
-            raise WorkflowRuntimeError("Workflow instance no longer available")
-        state = init_state or BrokerState.from_workflow(wf)
-        runner = _ControlLoopRunner(wf, plugin, context, step_workers, state)
-        return await runner.run(start_event=start_event)
-
-    return control_loop
+    # Prefer run-scoped context if available (set by broker)
+    current = workflow_registry.get_run(run_id)
+    if current is None:
+        raise WorkflowRuntimeError("Run context not found for control loop")
+    state = init_state or BrokerState.from_workflow(current.workflow)
+    runner = _ControlLoopRunner(
+        current.workflow, current.plugin, current.context, current.steps, state
+    )
+    return await runner.run(start_event=start_event)
 
 
 def rebuild_state_from_ticks(

src/workflows/runtime/types/plugin.py

@@ -130,8 +130,5 @@ def __call__(
         self,
         start_event: Event | None,
         init_state: BrokerState | None,
-        # these will likely be refactored out from the control loop function in the future.
-        plugin: WorkflowRuntime,
-        context: Context,
-        step_workers: dict[str, StepWorkerFunction],
+        run_id: str,
     ) -> Coroutine[None, None, StopEvent]: ...

src/workflows/runtime/workflow_registry.py

@@ -1,13 +1,19 @@
 from threading import Lock
 from weakref import WeakKeyDictionary
+from dataclasses import dataclass
+from typing import Optional
 from workflows.runtime.types.plugin import (
     ControlLoopFunction,
     Plugin,
     RegisteredWorkflow,
+    WorkflowRuntime,
 )
 from workflows.workflow import Workflow
 from workflows.runtime.types.step_function import StepWorkerFunction
-from workflows.decorators import R
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from workflows.context.context import Context
 
 
 class WorkflowPluginRegistry:
@@ -22,13 +28,14 @@ def __init__(self) -> None:
             Workflow, dict[type[Plugin], RegisteredWorkflow]
         ] = WeakKeyDictionary()
         self.lock = Lock()
+        self.run_contexts: dict[str, RegisteredRunContext] = {}
 
     def get_registered_workflow(
         self,
         workflow: Workflow,
         plugin: Plugin,
         workflow_function: ControlLoopFunction,
-        steps: dict[str, StepWorkerFunction[R]],
+        steps: dict[str, StepWorkerFunction],
     ) -> RegisteredWorkflow:
         plugin_type = type(plugin)
 
@@ -51,5 +58,38 @@ def get_registered_workflow(
             plugin_map[plugin_type] = registered_workflow
             return registered_workflow
 
+    def register_run(
+        self,
+        run_id: str,
+        workflow: Workflow,
+        plugin: WorkflowRuntime,
+        context: "Context",
+        steps: dict[str, StepWorkerFunction],
+    ) -> None:
+        self.run_contexts[run_id] = RegisteredRunContext(
+            run_id=run_id,
+            workflow=workflow,
+            plugin=plugin,
+            context=context,
+            steps=steps,
+        )
+
+    def get_run(self, run_id: str) -> Optional["RegisteredRunContext"]:
+        return self.run_contexts.get(run_id)
+
+    def delete_run(self, run_id: str) -> None:
+        self.run_contexts.pop(run_id, None)
+
 
 workflow_registry = WorkflowPluginRegistry()
+
+
+@dataclass
+class RegisteredRunContext:
+    run_id: str
+    workflow: Workflow
+    plugin: WorkflowRuntime
+    context: "Context"
+    steps: dict[str, StepWorkerFunction]
+
+