Add dbos plugin

Author: adrianlyjak

pyproject.toml

@@ -27,12 +27,13 @@ dependencies = [
   "eval-type-backport>=0.2.2 ; python_full_version < '3.10'",
   "llama-index-instrumentation>=0.1.0",
   "pydantic>=2.11.5",
-  "typing-extensions>=4.6.0"
+  "typing-extensions>=4.6.0",
 ]
 
 [project.optional-dependencies]
 server = ["starlette>=0.39.0", "uvicorn>=0.32.0"]
 client = ["httpx>=0.28.1,<1"]
+dbos = ["dbos>=1.14.0"]
 
 [tool.basedpyright]
 typeCheckingMode = "standard"

src/workflows/plugins/dbos.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2025 LlamaIndex Inc.
+
+from __future__ import annotations
+
+import time
+from typing import AsyncGenerator
+
+from dbos import DBOS, SetWorkflowID  # required extra, import must succeed
+
+from workflows.events import Event, StopEvent
+from workflows.runtime.types.plugin import (
+    ControlLoopFunction,
+    Plugin,
+    WorkflowRuntime,
+    RegisteredWorkflow,
+)
+from workflows.runtime.types.internal_state import BrokerState
+from workflows.runtime.types.step_function import StepWorkerFunction
+from workflows.runtime.types.ticks import WorkflowTick
+
+from workflows.workflow import Workflow
+
+
+@DBOS.step()
+async def _durable_time() -> float:
+    return time.time()
+
+
+class DBOSRuntime:
+    def register(
+        self,
+        workflow: Workflow,
+        workflow_function: ControlLoopFunction,
+        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.workflow()
+        async def _dbos_control_loop(
+            start_event: Event | None,
+            init_state: BrokerState | None,
+            run_id: str,
+        ) -> StopEvent:
+            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,
+            run_id: str,
+        ) -> StopEvent:
+            # Call the DBOS workflow directly; DBOS will orchestrate execution
+            return await _dbos_control_loop(start_event, init_state, run_id)
+
+        return RegisteredWorkflow(workflow_function=_dbos_control_loop, steps=steps)
+
+    def new_runtime(self, run_id: str) -> WorkflowRuntime:
+        runtime: WorkflowRuntime = DBOSWorkflowRuntime(run_id)
+        return runtime
+
+
+dbos_runtime: Plugin = DBOSRuntime()
+
+
+class DBOSWorkflowRuntime:
+    """
+    Workflow runtime backed by asyncio mailboxes, with durable timing via DBOS when available.
+
+    - send_event/wait_receive implement the tick mailbox used by the control loop
+    - write_to_event_stream/stream_published_events expose published events to callers
+    - get_now returns a stable value on first call within a run (durable if DBOS is installed)
+    - sleep uses DBOS durable sleep when available, otherwise asyncio.sleep
+    - on_tick/replay provide a lightweight snapshot for debug/replay via the broker
+    """
+
+    def __init__(
+        self,
+        run_id: str,
+    ) -> None:
+        self.run_id = run_id
+
+    # Mailbox used by control loop and broker
+    async def wait_receive(self) -> WorkflowTick:
+        # Receive next tick via DBOS durable notification
+        tick = await DBOS.recv_async()
+        return tick  # type: ignore[return-value]
+
+    async def send_event(self, tick: WorkflowTick) -> None:
+        await DBOS.send_async(self.run_id, tick)
+
+    # Event stream used by handlers/observers
+    async def write_to_event_stream(self, event: Event) -> None:
+        await DBOS.write_stream_async("published_events", event)
+
+    async def stream_published_events(self) -> AsyncGenerator[Event, None]:
+        async for event in DBOS.read_stream_async(self.run_id, "published_events"):
+            yield event
+
+    # Timing utilities
+    async def get_now(self) -> float:
+        return await _durable_time()
+
+    async def sleep(self, seconds: float) -> None:
+        await DBOS.sleep_async(seconds)
+
+    async def close(self) -> None:
+        pass

tests/plugins/test_dbos_plugin.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import Generator, Any
+from pathlib import Path
+
+import pytest
+
+# Require Python 3.10+ for dbos typing
+# if sys.version_info < (3, 10):  # pragma: no cover - environment guard
+#     pytest.skip("Requires Python 3.10 or higher", allow_module_level=True)
+
+from dbos import DBOS, DBOSConfig
+from workflows.plugins.dbos import dbos_runtime
+
+from workflows.workflow import Workflow
+from workflows.decorators import step
+from workflows.events import StartEvent, StopEvent
+from workflows.context.context import Context
+
+
+@pytest.fixture()
+def dbos(tmp_path: Path) -> Generator[None, Any, None]:
+    # Use a file-based SQLite DB so the schema persists across connections/threads
+    db_file: Path = tmp_path / "dbos_test.sqlite3"
+    # Allow usage across threads in tests
+    system_db_url: str = f"sqlite+pysqlite:///{db_file}?check_same_thread=false"
+
+    config: DBOSConfig = {
+        "name": "workflows-py-dbostest",
+        "system_database_url": system_db_url,
+        "run_admin_server": False,
+    }
+    # DBOS.reset_system_database()
+    DBOS(config=config)
+    DBOS.launch()
+    try:
+        yield None
+    finally:
+        DBOS.destroy()
+
+
+class SimpleWorkflow(Workflow):
+    @step
+    async def start(self, ev: StartEvent) -> StopEvent:
+        return StopEvent(result="ok")
+
+
+@pytest.mark.asyncio
+async def test_dbos_plugin_simple_run(dbos: None) -> None:
+    wf = SimpleWorkflow()
+    ctx: Context = Context(wf, plugin=dbos_runtime)
+    handler = wf.run(ctx=ctx)
+    result = await handler
+    assert result == "ok"