jank

Author: adrianlyjak

packages/llama-index-workflows/src/workflows/context/context.py

@@ -283,6 +283,7 @@ def _workflow_run(
                 start_event=start_event,
                 serialized_state=pre.serialized_state,
                 serializer=pre.serializer,
+                adapter_state=pre.init_snapshot.adapter_state,
             )
 
             # TODO(v3): Remove mutation. Handler will just be the external face.

packages/llama-index-workflows/src/workflows/context/context_types.py

@@ -131,6 +131,10 @@ class SerializedContext(BaseModel):
     # Maps step_name -> SerializedStepWorkerState
     workers: dict[str, SerializedStepWorkerState] = Field(default_factory=dict)
 
+    # Opaque adapter state for runtime-specific pending data (e.g., receive_queue contents).
+    # Each adapter defines its own structure. None means no pending state.
+    adapter_state: Optional[dict[str, Any]] = Field(default=None)
+
     @staticmethod
     def from_v0(v0: SerializedContextV0) -> "SerializedContext":
         """Convert V0 format to current format.

packages/llama-index-workflows/src/workflows/context/external_context.py

@@ -165,6 +165,14 @@ def to_dict(self, serializer: BaseSerializer | None = None) -> dict[str, Any]:
 
         context = broker_state.to_serialized(active_serializer)
         context.state = state_data
+
+        # Get pending adapter state (e.g., events in receive_queue)
+        snapshottable = as_snapshottable_adapter(self._external_adapter)
+        if snapshottable is not None:
+            adapter_state = snapshottable.get_pending_state()
+            if adapter_state is not None:
+                context.adapter_state = adapter_state
+
         return context.model_dump(mode="python")
 
     def cancel(self) -> None:

packages/llama-index-workflows/src/workflows/plugins/basic.py

@@ -33,7 +33,7 @@
     as_step_worker_functions,
     create_workflow_run_function,
 )
-from workflows.runtime.types.ticks import WorkflowTick
+from workflows.runtime.types.ticks import WorkflowTick, tick_from_dict, tick_to_dict
 from workflows.workflow import Workflow
 
 
@@ -53,11 +53,13 @@ def __init__(
         run_id: str,
         init_state: BrokerState,
         state_store: InMemoryStateStore[Any] | None = None,
+        serializer: BaseSerializer | None = None,
     ):
         self.run_id = run_id
         self.init_state = init_state
         self.ticks: list[WorkflowTick] = []
         self.state_store = state_store
+        self.serializer = serializer or JsonSerializer()
 
     # created lazily via cached_property for Python 3.14+ compatibility (they require a running event loop)
     @functools.cached_property
@@ -193,6 +195,44 @@ def abort(self) -> None:
     def init_state(self) -> BrokerState:
         return self._queues.init_state
 
+    def get_pending_state(self) -> dict[str, Any] | None:
+        """Snapshot receive_queue contents for serialization.
+
+        Returns serialized pending ticks that haven't been processed yet.
+        These are events sent via ctx.send_event() that are still in the queue.
+
+        Note: TickCancelRun is filtered out since cancel is a transient operation
+        that shouldn't survive serialization/resume cycles.
+        """
+        from workflows.runtime.types.ticks import TickCancelRun
+
+        # Access the internal deque of asyncio.Queue
+        # This is a private attribute but stable across Python versions
+        queue = self._queues.receive_queue
+        pending_ticks = [
+            tick
+            for tick in queue._queue  # type: ignore[attr-defined]
+            if not isinstance(tick, TickCancelRun)
+        ]
+        if not pending_ticks:
+            return None
+        return {
+            "pending_ticks": [
+                tick_to_dict(tick, self._queues.serializer) for tick in pending_ticks
+            ]
+        }
+
+    def restore_pending_state(self, state: dict[str, Any]) -> None:
+        """Restore pending ticks to receive_queue on resume.
+
+        Re-populates the receive_queue with ticks that were pending
+        at serialization time.
+        """
+        pending_ticks_data = state.get("pending_ticks", [])
+        for tick_data in pending_ticks_data:
+            tick = tick_from_dict(tick_data, self._queues.serializer)
+            self._queues.receive_queue.put_nowait(tick)
+
 
 class BasicRuntime(Runtime):
     """Default asyncio-based runtime with no durability."""
@@ -251,6 +291,7 @@ def run_workflow(
         start_event: StartEvent | None = None,
         serialized_state: dict[str, Any] | None = None,
         serializer: BaseSerializer | None = None,
+        adapter_state: dict[str, Any] | None = None,
     ) -> ExternalRunAdapter:
         """Set up a workflow run. Currently only creates state store.
 
@@ -276,6 +317,12 @@ def run_workflow(
         # might want to lock this better. Unlikely race condition if you spam with the same run_id.
         queues = self._get_or_create_queues(run_id, init_state)
         queues.state_store = state_store
+        queues.serializer = active_serializer
+
+        # Restore adapter state (pending ticks) if resuming
+        if adapter_state is not None:
+            external_adapter = ExternalAsyncioAdapter(queues)
+            external_adapter.restore_pending_state(adapter_state)
 
         async def run_with_concurrency_limit() -> StopEvent:
             # Capture strong reference to queues for the task's lifetime,

packages/llama-index-workflows/src/workflows/plugins/dbos.py

@@ -144,6 +144,7 @@ def run_workflow(
         start_event: StartEvent | None = None,
         serialized_state: dict[str, Any] | None = None,
         serializer: BaseSerializer | None = None,
+        adapter_state: dict[str, Any] | None = None,
     ) -> ExternalRunAdapter:
         """Set up a workflow run. Currently only creates state store.
 

packages/llama-index-workflows/src/workflows/runtime/control_loop.py

@@ -313,21 +313,17 @@ async def run(
 
         # Resume any in-progress work
         self.state, commands = rewind_in_progress(self.state, now)
-        print(f"DEBUG: rewind_in_progress returned {len(commands)} commands")
         for command in commands:
-            print(f"DEBUG: processing command {type(command).__name__}")
             try:
                 await self.process_command(command)
             except Exception:
                 await self.cleanup_tasks()
                 raise
-        print(f"DEBUG: worker_tasks after rewind: {len(self.worker_tasks)}")
 
         # Initialize pull task (single-iteration)
         pull_task: asyncio.Task[WorkflowTick | None] | None = asyncio.create_task(
             _single_pull(self.adapter)
         )
-        print(f"DEBUG: pull_task created")
 
         # Main event loop
         try:
@@ -350,10 +346,6 @@ async def run(
                 # Drain and process buffered ticks first (from rehydration, queue_tick, etc.)
                 while self.tick_buffer:
                     tick = self.tick_buffer.pop(0)
-                    tick_detail = ""
-                    if hasattr(tick, 'event'):
-                        tick_detail = f" ({type(tick.event).__name__})"  # type: ignore
-                    print(f"DEBUG: processing tick from buffer: {type(tick).__name__}{tick_detail}")
                     result = await self._process_tick(tick)
                     if result is not None:
                         return result
@@ -368,15 +360,13 @@ async def run(
 
                 # Gather all tasks to wait on
                 all_tasks: set[asyncio.Task[Any]] = {pull_task, *self.worker_tasks}
-                print(f"DEBUG: waiting on {len(all_tasks)} tasks (workers={len(self.worker_tasks)}, timeout={timeout})")
 
                 # Wait for first completion (with timeout for scheduled wakeups)
                 done, _ = await asyncio.wait(
                     all_tasks,
                     timeout=timeout,
                     return_when=asyncio.FIRST_COMPLETED,
                 )
-                print(f"DEBUG: asyncio.wait returned, done={len(done)}")
 
                 if not done:
                     # Timeout - process scheduled ticks
@@ -400,7 +390,6 @@ async def run(
                         except Exception:
                             logger.exception("Worker task failed unexpectedly")
                             continue
-                        print(f"DEBUG: worker completed: {tick_result.step_name}")
                         result = await self._process_tick(tick_result)
                         if result is not None:
                             return result
@@ -418,10 +407,6 @@ async def run(
                         # Respawn before processing to keep pulling during tick processing
                         pull_task = asyncio.create_task(_single_pull(self.adapter))
                         if pull_tick is not None:
-                            tick_detail = ""
-                            if hasattr(pull_tick, 'event'):
-                                tick_detail = f" ({type(pull_tick.event).__name__})"  # type: ignore
-                            print(f"DEBUG: pull received: {type(pull_tick).__name__}{tick_detail}")
                             result = await self._process_tick(pull_tick)
                             if result is not None:
                                 return result

packages/llama-index-workflows/src/workflows/runtime/types/plugin.py

@@ -328,6 +328,7 @@ def run_workflow(
         start_event: StartEvent | None = None,
         serialized_state: dict[str, Any] | None = None,
         serializer: "BaseSerializer | None" = None,
+        adapter_state: dict[str, Any] | None = None,
     ) -> ExternalRunAdapter:
         """
         Launch a workflow run.
@@ -342,6 +343,7 @@ def run_workflow(
             start_event: Optional start event to begin the workflow.
             serialized_state: Serialized state store data to restore from.
             serializer: Serializer to use for deserializing state.
+            adapter_state: Opaque adapter state to restore (e.g., pending events).
         """
         ...
 
@@ -462,6 +464,32 @@ def replay(self) -> list[WorkflowTick]:
         """
         ...
 
+    def get_pending_state(self) -> dict[str, Any] | None:
+        """
+        Return pending adapter state for serialization.
+
+        Called during context serialization to capture any internal state
+        (like pending events in a receive queue) that should be persisted.
+        This state will be restored via restore_pending_state() on resume.
+
+        Returns:
+            dict with adapter-specific state, or None if no state to persist.
+            The format is opaque - each adapter defines its own structure.
+        """
+        return None
+
+    def restore_pending_state(self, state: dict[str, Any]) -> None:
+        """
+        Restore pending adapter state on resume.
+
+        Called when resuming a workflow to restore any internal state
+        that was captured by get_pending_state() during serialization.
+
+        Args:
+            state: The adapter-specific state dict from get_pending_state()
+        """
+        pass
+
 
 def as_snapshottable_adapter(
     adapter: ExternalRunAdapter | InternalRunAdapter,

packages/llama-index-workflows/src/workflows/runtime/types/ticks.py

@@ -17,12 +17,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Generic, Union
+from typing import TYPE_CHECKING, Generic, Union
 
 from workflows.decorators import R
 from workflows.events import Event
 from workflows.runtime.types.results import StepFunctionResult
 
+if TYPE_CHECKING:
+    from workflows.context.serializers import BaseSerializer
+
 
 @dataclass(frozen=True)
 class TickStepResult(Generic[R]):
@@ -68,3 +71,72 @@ class TickTimeout:
 WorkflowTick = Union[
     TickStepResult[R], TickAddEvent, TickCancelRun, TickPublishEvent, TickTimeout
 ]
+
+
+def tick_to_dict(tick: WorkflowTick, serializer: "BaseSerializer") -> dict:
+    """Serialize a WorkflowTick to a dict for persistence.
+
+    Args:
+        tick: The tick to serialize
+        serializer: Serializer for event payloads
+
+    Returns:
+        A dict representation that can be JSON-serialized
+    """
+    if isinstance(tick, TickAddEvent):
+        return {
+            "type": "TickAddEvent",
+            "event": serializer.serialize(tick.event),
+            "step_name": tick.step_name,
+            "attempts": tick.attempts,
+            "first_attempt_at": tick.first_attempt_at,
+        }
+    elif isinstance(tick, TickCancelRun):
+        return {"type": "TickCancelRun"}
+    elif isinstance(tick, TickPublishEvent):
+        return {
+            "type": "TickPublishEvent",
+            "event": serializer.serialize(tick.event),
+        }
+    elif isinstance(tick, TickTimeout):
+        return {
+            "type": "TickTimeout",
+            "timeout": tick.timeout,
+        }
+    elif isinstance(tick, TickStepResult):
+        # TickStepResult shouldn't be in the pending queue, but handle it for completeness
+        raise ValueError(
+            "TickStepResult cannot be serialized - it should not be in pending state"
+        )
+    else:
+        raise ValueError(f"Unknown tick type: {type(tick)}")
+
+
+def tick_from_dict(data: dict, serializer: "BaseSerializer") -> WorkflowTick:
+    """Deserialize a WorkflowTick from a dict.
+
+    Args:
+        data: The dict representation from tick_to_dict
+        serializer: Serializer for event payloads
+
+    Returns:
+        The deserialized WorkflowTick
+    """
+    tick_type = data.get("type")
+    if tick_type == "TickAddEvent":
+        return TickAddEvent(
+            event=serializer.deserialize(data["event"]),
+            step_name=data.get("step_name"),
+            attempts=data.get("attempts"),
+            first_attempt_at=data.get("first_attempt_at"),
+        )
+    elif tick_type == "TickCancelRun":
+        return TickCancelRun()
+    elif tick_type == "TickPublishEvent":
+        return TickPublishEvent(
+            event=serializer.deserialize(data["event"]),
+        )
+    elif tick_type == "TickTimeout":
+        return TickTimeout(timeout=data["timeout"])
+    else:
+        raise ValueError(f"Unknown tick type: {tick_type}")

packages/llama-index-workflows/tests/runtime/conftest.py

@@ -70,6 +70,7 @@ def run_workflow(
         start_event: StartEvent | None = None,
         serialized_state: dict[str, Any] | None = None,
         serializer: "BaseSerializer | None" = None,
+        adapter_state: dict[str, Any] | None = None,
     ) -> ExternalRunAdapter:
         self._current_run_id = run_id
         return self.get_external_adapter(run_id)