pyworkflow-engine 0.2.6__tar.gz → 0.2.7__tar.gz

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.2.6
+Version: 0.2.7
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.2.6"
+version = "0.2.7"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.2.6"
+__version__ = "0.2.7"
 
 # Configuration
 from pyworkflow.config import (

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/celery/app.py

@@ -335,6 +335,11 @@ def create_celery_app(
                 "schedule": timedelta(hours=24),
                 "options": {"queue": "pyworkflow.default"},
             },
+            "pyworkflow-streams-drain-scheduled": {
+                "task": "pyworkflow.streams.drain_scheduled_signals",
+                "schedule": timedelta(seconds=2),
+                "options": {"queue": "pyworkflow.default"},
+            },
         },
         # Logging
         "worker_log_format": "[%(asctime)s: %(levelname)s/%(processName)s] %(message)s",

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/celery/tasks.py

@@ -2840,3 +2840,103 @@ def run_data_retention_task(self: SingletonWorkflowTask) -> dict[str, Any]:
     count = run_async(backend.delete_old_runs(cutoff))
     logger.info("Data retention complete: deleted {} runs", count)
     return {"deleted": count, "skipped": False}
+
+
+# =============================================================================
+# Stream-only celery tasks (private to pyworkflow.streams)
+# =============================================================================
+
+
+@celery_app.task(
+    name="pyworkflow.streams.dispatch_signal",
+    queue="pyworkflow.default",
+)
+def dispatch_stream_signal_task(
+    signal_id: str,
+    stream_id: str,
+    signal_type: str,
+    payload: dict[str, Any],
+    sequence: int,
+    source_run_id: str | None,
+    stream_run_id: str | None,
+    metadata: dict[str, Any] | None,
+    storage_config: dict[str, Any] | None = None,
+) -> None:
+    """Dispatch a previously persisted signal to its subscribed stream steps.
+
+    Reconstructs the Signal from inline kwargs (avoids a DB round-trip),
+    runs the dispatcher logic on the worker, and after step execution
+    completes calls _maybe_resume_stream_parent so the parent @workflow
+    can be resumed via resume_hook() if the aggregate is now terminal.
+    """
+    from pyworkflow.celery.app import _configure_worker_logging
+
+    _configure_worker_logging()
+
+    storage = _get_storage_backend(storage_config)
+
+    from pyworkflow.streams.dispatcher import dispatch_signal
+    from pyworkflow.streams.signal import Signal
+
+    signal = Signal(
+        signal_id=signal_id,
+        stream_id=stream_id,
+        signal_type=signal_type,
+        payload=payload or {},
+        sequence=sequence,
+        source_run_id=source_run_id,
+        stream_run_id=stream_run_id,
+        metadata=metadata or {},
+    )
+
+    run_async(dispatch_signal(signal, storage))
+
+
+@celery_app.task(
+    name="pyworkflow.streams.drain_scheduled_signals",
+    base=SingletonWorkflowTask,
+    bind=True,
+    queue="pyworkflow.default",
+    release_lock_on_failure=True,
+    lock_expiry=300,
+)
+def drain_scheduled_signals_task(self: SingletonWorkflowTask) -> dict[str, Any]:
+    """Periodic task: deliver any due rows from the ``scheduled_signals``
+    table by calling :func:`pyworkflow.streams.emit.emit`. Replaces the
+    in-process poller that previously ran inside ``run_stream_workflow``."""
+    from pyworkflow.celery.app import _configure_worker_logging
+    from pyworkflow.config import get_config
+
+    _configure_worker_logging()
+
+    config = get_config()
+    storage = config.storage
+    if storage is None:
+        return {"delivered": 0, "skipped": True}
+
+    async def _drain() -> int:
+        from pyworkflow.streams.emit import emit
+
+        try:
+            due = await storage.fetch_due_scheduled_signals(datetime.now(UTC), limit=100)
+        except NotImplementedError:
+            return 0
+        delivered = 0
+        for row in due or []:
+            try:
+                await emit(
+                    stream_id=row["stream_id"],
+                    signal_type=row["signal_type"],
+                    payload=row.get("payload") or {},
+                    storage=storage,
+                    stream_run_id=row.get("stream_run_id"),
+                    metadata=row.get("metadata") or {},
+                )
+                await storage.mark_scheduled_signal_delivered(row["id"])
+                delivered += 1
+            except Exception as e:  # noqa: BLE001
+                logger.error(f"[drain_scheduled_signals] failed to deliver {row.get('id')}: {e}")
+        return delivered
+
+    count = run_async(_drain())
+    return {"delivered": count, "skipped": False}

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/storage/base.py

@@ -7,6 +7,7 @@ across different backends (File, Redis, SQLite, PostgreSQL).
 
 from abc import ABC, abstractmethod
 from datetime import datetime
+from typing import Any
 
 from pyworkflow.engine.events import Event
 from pyworkflow.storage.schemas import (
@@ -841,6 +842,7 @@ class StorageBackend(ABC):
         stream_id: str,
         step_run_id: str,
         signal_types: list[str],
+        stream_run_id: str | None = None,
     ) -> None:
         """
         Register a stream step's subscription to signal types.
@@ -859,6 +861,7 @@ class StorageBackend(ABC):
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """
         Get step_run_ids waiting for a specific signal type on a stream.
@@ -879,6 +882,7 @@ class StorageBackend(ABC):
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """
         Get ALL subscriptions for a signal type on a stream, regardless of status.
@@ -897,6 +901,40 @@ class StorageBackend(ABC):
         """
         raise NotImplementedError("This storage backend does not support streams")
 
+    async def set_subscription_result(
+        self,
+        stream_id: str,
+        step_run_id: str,
+        result: Any,
+    ) -> None:
+        """Persist a stream step's result payload on its subscription row.
+
+        Returned to the parent ``@workflow`` via
+        ``StreamWorkflowResult.step_results`` so it doesn't have to fish
+        through checkpoints to read step output.
+        """
+        raise NotImplementedError("This storage backend does not support streams")
+
+    async def set_stream_parent_link(
+        self,
+        stream_id: str,
+        stream_run_id: str,
+        parent_run_id: str,
+        parent_hook_token: str,
+    ) -> None:
+        """Record (parent_run_id, parent_hook_token) for a stream run so the
+        stream-step dispatcher can resume the parent @workflow via resume_hook()
+        once the stream reaches a terminal aggregate state."""
+        raise NotImplementedError("This storage backend does not support streams")
+
+    async def get_stream_parent_link(
+        self,
+        stream_id: str,
+        stream_run_id: str,
+    ) -> tuple[str, str] | None:
+        """Return the (parent_run_id, parent_hook_token) recorded for a stream run."""
+        raise NotImplementedError("This storage backend does not support streams")
+
     async def update_subscription_status(
         self,
         stream_id: str,
@@ -909,13 +947,40 @@ class StorageBackend(ABC):
         Args:
             stream_id: Stream identifier
             step_run_id: The stream step's run ID
-            status: New status ("waiting", "running", "completed")
+            status: New status. One of:
+                - "waiting": ready for next signal
+                - "running": currently executing lifecycle
+                - "suspended": alive but waiting for an external resume condition
+                  (e.g. HITL). Contributes "suspended" to the stream-workflow
+                  aggregate state.
+                - "terminated": permanently done; dispatcher will skip it.
+                - "completed": legacy alias for "terminated".
 
         Raises:
             NotImplementedError: If backend doesn't support streams
         """
         raise NotImplementedError("This storage backend does not support streams")
 
+    async def get_subscription_states(
+        self,
+        stream_id: str,
+        stream_run_id: str | None = None,
+    ) -> list[dict]:
+        """
+        Return all subscriptions and their statuses for a stream.
+
+        Used by the stream-workflow runtime to compute the aggregate state
+        across all registered ``@stream_step``s.
+
+        Args:
+            stream_id: Stream identifier
+            stream_run_id: Optional stream run scope
+
+        Returns:
+            List of dicts with at least ``step_run_id`` and ``status`` keys.
+        """
+        raise NotImplementedError("This storage backend does not support streams")
+
     async def acknowledge_signal(
         self,
         signal_id: str,
@@ -953,6 +1018,41 @@ class StorageBackend(ABC):
         """
         raise NotImplementedError("This storage backend does not support streams")
 
+    # Scheduled Signal Operations (for schedule_signal primitive)
+
+    async def schedule_signal(
+        self,
+        *,
+        stream_id: str,
+        signal_type: str,
+        payload: dict,
+        due_at: "datetime",  # noqa: F821  (forward ref; datetime imported by backends)
+        stream_run_id: str | None = None,
+        metadata: dict | None = None,
+    ) -> str:
+        """
+        Persist a signal to be emitted at ``due_at``.
+
+        Used by the ``schedule_signal`` streams primitive for time-delayed
+        signal delivery (e.g. supervisor sleep/wakeup).
+
+        Returns:
+            Opaque scheduled-signal ID.
+        """
+        raise NotImplementedError("This storage backend does not support scheduled signals")
+
+    async def fetch_due_scheduled_signals(
+        self,
+        now: "datetime",  # noqa: F821
+        limit: int = 100,
+    ) -> list[dict]:
+        """Return scheduled signals whose ``due_at <= now`` and not yet delivered."""
+        raise NotImplementedError("This storage backend does not support scheduled signals")
+
+    async def mark_scheduled_signal_delivered(self, sched_id: str) -> None:
+        """Mark a scheduled signal as delivered so it is not emitted twice."""
+        raise NotImplementedError("This storage backend does not support scheduled signals")
+
     # Checkpoint Operations
 
     async def save_checkpoint(self, step_run_id: str, data: dict) -> None:

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/storage/cassandra.py

@@ -2101,6 +2101,7 @@ class CassandraStorageBackend(StorageBackend):
         stream_id: str,
         step_run_id: str,
         signal_types: list[str],
+        stream_run_id: str | None = None,
     ) -> None:
         """Register a stream step's subscription to signal types."""
         session = self._ensure_connected()
@@ -2119,6 +2120,7 @@ class CassandraStorageBackend(StorageBackend):
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """Get step_run_ids waiting for a specific signal type on a stream."""
         session = self._ensure_connected()

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/storage/citus.py

@@ -110,6 +110,95 @@ class CitusMigrationRunner(PostgresMigrationRunner):
                 # Re-distribute the new tables
                 await conn.execute("SELECT create_distributed_table('signals', 'stream_run_id')")
                 await conn.execute("SELECT create_reference_table('signal_acknowledgments')")
+            elif migration.version == 4:
+                # V4: Drop FK constraints referencing streams table.
+                # Citus never created these FKs (cross-shard FKs unsupported),
+                # so this is a no-op but we still record the version.
+                pass
+            elif migration.version == 5:
+                # V5: Add stream_run_id to stream_subscriptions + scheduled_signals table
+                await conn.execute("""
+                    DO $$
+                    BEGIN
+                        IF EXISTS (
+                            SELECT FROM information_schema.tables
+                            WHERE table_name = 'stream_subscriptions'
+                        ) AND NOT EXISTS (
+                            SELECT 1 FROM information_schema.columns
+                            WHERE table_name = 'stream_subscriptions'
+                              AND column_name = 'stream_run_id'
+                        ) THEN
+                            ALTER TABLE stream_subscriptions ADD COLUMN stream_run_id TEXT NULL;
+                        END IF;
+                    END $$
+                """)
+                await conn.execute(
+                    "CREATE INDEX IF NOT EXISTS idx_subscriptions_stream_run "
+                    "ON stream_subscriptions(stream_id, stream_run_id)"
+                )
+                await conn.execute("""
+                    CREATE TABLE IF NOT EXISTS scheduled_signals (
+                        id TEXT PRIMARY KEY,
+                        stream_id TEXT NOT NULL,
+                        signal_type TEXT NOT NULL,
+                        payload JSONB NOT NULL DEFAULT '{}',
+                        due_at TIMESTAMPTZ NOT NULL,
+                        stream_run_id TEXT,
+                        metadata JSONB DEFAULT '{}',
+                        delivered BOOLEAN NOT NULL DEFAULT FALSE,
+                        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+                    )
+                """)
+                await conn.execute(
+                    "CREATE INDEX IF NOT EXISTS idx_scheduled_signals_due "
+                    "ON scheduled_signals(delivered, due_at)"
+                )
+                # Distribute as reference table if not already distributed
+                already = await conn.fetchval("""
+                    SELECT EXISTS (
+                        SELECT 1 FROM pg_dist_partition dp
+                        JOIN pg_class c ON dp.logicalrelid = c.oid
+                        WHERE c.relname = 'scheduled_signals'
+                    )
+                """)
+                if not already:
+                    await conn.execute("SELECT create_reference_table('scheduled_signals')")
+            elif migration.version == 6:
+                # V6: Add parent_run_id + parent_hook_token to stream_subscriptions
+                await conn.execute("""
+                    DO $$
+                    BEGIN
+                        IF EXISTS (
+                            SELECT FROM information_schema.tables
+                            WHERE table_name = 'stream_subscriptions'
+                        ) AND NOT EXISTS (
+                            SELECT 1 FROM information_schema.columns
+                            WHERE table_name = 'stream_subscriptions'
+                              AND column_name = 'parent_run_id'
+                        ) THEN
+                            ALTER TABLE stream_subscriptions
+                                ADD COLUMN parent_run_id TEXT NULL,
+                                ADD COLUMN parent_hook_token TEXT NULL;
+                        END IF;
+                    END $$
+                """)
+            elif migration.version == 7:
+                await conn.execute("""
+                    DO $$
+                    BEGIN
+                        IF EXISTS (
+                            SELECT FROM information_schema.tables
+                            WHERE table_name = 'stream_subscriptions'
+                        ) AND NOT EXISTS (
+                            SELECT 1 FROM information_schema.columns
+                            WHERE table_name = 'stream_subscriptions'
+                              AND column_name = 'result'
+                        ) THEN
+                            ALTER TABLE stream_subscriptions
+                                ADD COLUMN result JSONB NULL;
+                        END IF;
+                    END $$
+                """)
             elif migration.up_func:
                 await migration.up_func(conn)
             elif migration.up_sql and migration.up_sql != "SELECT 1":
@@ -410,6 +499,10 @@ class CitusStorageBackend(PostgresStorageBackend):
                     step_run_id TEXT NOT NULL,
                     signal_types JSONB NOT NULL DEFAULT '[]',
                     status TEXT NOT NULL DEFAULT 'waiting',
+                    stream_run_id TEXT NULL,
+                    parent_run_id TEXT NULL,
+                    parent_hook_token TEXT NULL,
+                    result JSONB NULL,
                     created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                     PRIMARY KEY (stream_id, step_run_id)
                 )
@@ -417,6 +510,29 @@ class CitusStorageBackend(PostgresStorageBackend):
             await conn.execute(
                 "CREATE INDEX IF NOT EXISTS idx_subscriptions_status ON stream_subscriptions(stream_id, status)"
             )
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_subscriptions_stream_run "
+                "ON stream_subscriptions(stream_id, stream_run_id)"
+            )
+
+            # scheduled_signals — reference table (polled globally by runtime)
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS scheduled_signals (
+                    id TEXT PRIMARY KEY,
+                    stream_id TEXT NOT NULL,
+                    signal_type TEXT NOT NULL,
+                    payload JSONB NOT NULL DEFAULT '{}',
+                    due_at TIMESTAMPTZ NOT NULL,
+                    stream_run_id TEXT,
+                    metadata JSONB DEFAULT '{}',
+                    delivered BOOLEAN NOT NULL DEFAULT FALSE,
+                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+                )
+            """)
+            await conn.execute(
+                "CREATE INDEX IF NOT EXISTS idx_scheduled_signals_due "
+                "ON scheduled_signals(delivered, due_at)"
+            )
 
             # signal_acknowledgments — reference table (needs cross-shard joins)
             # Differences vs postgres:
@@ -496,3 +612,7 @@ class CitusStorageBackend(PostgresStorageBackend):
             # signal_acknowledgments: reference table (needs cross-shard joins)
             if "signal_acknowledgments" not in distributed:
                 await conn.execute("SELECT create_reference_table('signal_acknowledgments')")
+
+            # scheduled_signals: reference table (polled globally by runtime)
+            if "scheduled_signals" not in distributed:
+                await conn.execute("SELECT create_reference_table('scheduled_signals')")

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/storage/dynamodb.py

@@ -1686,6 +1686,7 @@ class DynamoDBStorageBackend(StorageBackend):
         stream_id: str,
         step_run_id: str,
         signal_types: list[str],
+        stream_run_id: str | None = None,
     ) -> None:
         """Register a stream step's subscription to signal types."""
         async with self._get_client() as client:
@@ -1710,6 +1711,7 @@ class DynamoDBStorageBackend(StorageBackend):
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """Get step_run_ids waiting for a specific signal type on a stream."""
         async with self._get_client() as client:

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/storage/file.py

@@ -1187,6 +1187,7 @@ class FileStorageBackend(StorageBackend):
         stream_id: str,
         step_run_id: str,
         signal_types: list[str],
+        stream_run_id: str | None = None,
     ) -> None:
         """Register a stream step's subscription to signal types."""
         sub_file = self.subscriptions_dir / f"{stream_id}__{step_run_id}.json"
@@ -1210,6 +1211,7 @@ class FileStorageBackend(StorageBackend):
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """Get step_run_ids waiting for a specific signal type on a stream."""
 

{pyworkflow_engine-0.2.6 → pyworkflow_engine-0.2.7}/pyworkflow/storage/memory.py

@@ -746,6 +746,7 @@ class InMemoryStorageBackend(StorageBackend):
         stream_id: str,
         step_run_id: str,
         signal_types: list[str],
+        stream_run_id: str | None = None,
     ) -> None:
         """Register a stream step's subscription to signal types."""
         with self._lock:
@@ -755,34 +756,136 @@ class InMemoryStorageBackend(StorageBackend):
                 "step_run_id": step_run_id,
                 "signal_types": signal_types,
                 "status": "waiting",
+                "stream_run_id": stream_run_id,
                 "created_at": datetime.now(UTC).isoformat(),
             }
 
+    async def get_subscription_states(
+        self,
+        stream_id: str,
+        stream_run_id: str | None = None,
+    ) -> list[dict]:
+        """Return all subscriptions and their statuses for a stream run."""
+        with self._lock:
+            result = []
+            for (sid, _), sub in self._subscriptions.items():
+                if sid != stream_id:
+                    continue
+                if stream_run_id is not None and sub.get("stream_run_id") not in (
+                    stream_run_id,
+                    None,
+                ):
+                    continue
+                result.append(
+                    {
+                        "stream_id": sub["stream_id"],
+                        "step_run_id": sub["step_run_id"],
+                        "status": sub["status"],
+                        "stream_run_id": sub.get("stream_run_id"),
+                    }
+                )
+            return result
+
+    async def schedule_signal(
+        self,
+        *,
+        stream_id: str,
+        signal_type: str,
+        payload: dict,
+        due_at: datetime,
+        stream_run_id: str | None = None,
+        metadata: dict | None = None,
+    ) -> str:
+        """Persist a signal to be emitted at ``due_at``."""
+        import uuid as _uuid
+
+        with self._lock:
+            if not hasattr(self, "_scheduled_signals"):
+                self._scheduled_signals: dict[str, dict] = {}
+            sched_id = f"sched_{_uuid.uuid4().hex[:16]}"
+            self._scheduled_signals[sched_id] = {
+                "id": sched_id,
+                "stream_id": stream_id,
+                "signal_type": signal_type,
+                "payload": payload,
+                "due_at": due_at,
+                "stream_run_id": stream_run_id,
+                "metadata": metadata or {},
+                "delivered": False,
+            }
+            return sched_id
+
+    async def fetch_due_scheduled_signals(
+        self,
+        now: datetime,
+        limit: int = 100,
+    ) -> list[dict]:
+        """Return scheduled signals whose ``due_at <= now`` and not delivered."""
+        with self._lock:
+            if not hasattr(self, "_scheduled_signals"):
+                return []
+            due = []
+            for row in self._scheduled_signals.values():
+                if row["delivered"]:
+                    continue
+                if row["due_at"] <= now:
+                    due.append(dict(row))
+                    if len(due) >= limit:
+                        break
+            return due
+
+    async def mark_scheduled_signal_delivered(self, sched_id: str) -> None:
+        with self._lock:
+            if not hasattr(self, "_scheduled_signals"):
+                return
+            row = self._scheduled_signals.get(sched_id)
+            if row:
+                row["delivered"] = True
+
     async def get_waiting_steps(
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """Get step_run_ids waiting for a specific signal type on a stream."""
         with self._lock:
             result = []
             for (sid, _), sub in self._subscriptions.items():
-                if sid == stream_id and sub["status"] == "waiting":
-                    if signal_type in sub["signal_types"]:
-                        result.append(sub)
+                if sid != stream_id:
+                    continue
+                if sub["status"] not in ("waiting", "suspended"):
+                    continue
+                if signal_type not in sub["signal_types"]:
+                    continue
+                if stream_run_id is not None and sub.get("stream_run_id") not in (
+                    stream_run_id,
+                    None,
+                ):
+                    continue
+                result.append(sub)
             return result
 
     async def get_subscriptions_for_stream(
         self,
         stream_id: str,
         signal_type: str,
+        stream_run_id: str | None = None,
     ) -> list[dict]:
         """Get ALL subscriptions for a signal type, regardless of status."""
         with self._lock:
             result = []
             for (sid, _), sub in self._subscriptions.items():
-                if sid == stream_id and signal_type in sub["signal_types"]:
-                    result.append(sub)
+                if sid != stream_id:
+                    continue
+                if signal_type not in sub["signal_types"]:
+                    continue
+                if stream_run_id is not None and sub.get("stream_run_id") not in (
+                    stream_run_id,
+                    None,
+                ):
+                    continue
+                result.append(sub)
             return result
 
     async def update_subscription_status(