定时任务订阅模式

2026-05-13 14:59:27 +08:00
1 parent 839386f0b4
commit aa034f8669
@@ -0,0 +1,66 @@
+"""Publish/subscribe support for scheduled tasks.
+
+Revision ID: 20260513_02
+Revises: 20260513_01
+Create Date: 2026-05-13
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import inspect
+
+
+revision = "20260513_02"
+down_revision = "20260513_01"
+branch_labels = None
+depends_on = None
+
+
+def _has_column(table_name: str, column_name: str) -> bool:
+    inspector = inspect(op.get_bind())
+    return any(column["name"] == column_name for column in inspector.get_columns(table_name))
+
+
+def _has_table(table_name: str) -> bool:
+    inspector = inspect(op.get_bind())
+    return table_name in inspector.get_table_names()
+
+
+def upgrade() -> None:
+    if not _has_column("scheduled_tasks", "published"):
+        op.add_column(
+            "scheduled_tasks",
+            sa.Column("published", sa.Boolean(), nullable=False, server_default=sa.false()),
+        )
+        op.create_index("ix_scheduled_tasks_published", "scheduled_tasks", ["published"])
+    if not _has_column("scheduled_tasks", "published_at"):
+        op.add_column(
+            "scheduled_tasks",
+            sa.Column("published_at", sa.DateTime(timezone=True), nullable=True),
+        )
+
+    if not _has_table("scheduled_task_subscriptions"):
+        op.create_table(
+            "scheduled_task_subscriptions",
+            sa.Column("id", sa.String(64), primary_key=True),
+            sa.Column("task_id", sa.String(64), nullable=False),
+            sa.Column("user_id", sa.String(64), nullable=False),
+            sa.Column("notify_element", sa.Boolean(), nullable=False, server_default=sa.false()),
+            sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+            sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
+            sa.UniqueConstraint("task_id", "user_id", name="uq_scheduled_task_subscriptions_task_user"),
+        )
+        op.create_index("ix_scheduled_task_subscriptions_task_id", "scheduled_task_subscriptions", ["task_id"])
+        op.create_index("ix_scheduled_task_subscriptions_user_id", "scheduled_task_subscriptions", ["user_id"])
+
+
+def downgrade() -> None:
+    if _has_table("scheduled_task_subscriptions"):
+        op.drop_table("scheduled_task_subscriptions")
+    if _has_column("scheduled_tasks", "published_at"):
+        op.drop_column("scheduled_tasks", "published_at")
+    if _has_column("scheduled_tasks", "published"):
+        op.drop_index("ix_scheduled_tasks_published", table_name="scheduled_tasks")
+        op.drop_column("scheduled_tasks", "published")
@@ -19,7 +19,13 @@ from deerflow.persistence.agents.model import AgentRow
 from deerflow.persistence.feedback.model import FeedbackRow
 from deerflow.persistence.models.run_event import RunEventRow
 from deerflow.persistence.run.model import RunRow
-from deerflow.persistence.scheduled_tasks.model import ScheduledTaskDeliveryProfileRow, ScheduledTaskRow, ScheduledTaskRunRow, SchedulerThreadRow
+from deerflow.persistence.scheduled_tasks.model import (
+    ScheduledTaskDeliveryProfileRow,
+    ScheduledTaskRow,
+    ScheduledTaskRunRow,
+    ScheduledTaskSubscriptionRow,
+    SchedulerThreadRow,
+)
 from deerflow.persistence.thread_meta.model import ThreadMetaRow
 from deerflow.persistence.user.model import UserRow

@@ -31,6 +37,7 @@ __all__ = [
    "ScheduledTaskDeliveryProfileRow",
    "ScheduledTaskRow",
    "ScheduledTaskRunRow",
+    "ScheduledTaskSubscriptionRow",
    "SchedulerThreadRow",
    "ThreadMetaRow",
    "UserRow",
@@ -4,7 +4,13 @@ from __future__ import annotations

 from deerflow.persistence.scheduled_tasks.base import ScheduledTaskStore
 from deerflow.persistence.scheduled_tasks.memory import MemoryScheduledTaskStore
-from deerflow.persistence.scheduled_tasks.model import ScheduledTaskDeliveryProfileRow, ScheduledTaskRow, ScheduledTaskRunRow, SchedulerThreadRow
+from deerflow.persistence.scheduled_tasks.model import (
+    ScheduledTaskDeliveryProfileRow,
+    ScheduledTaskRow,
+    ScheduledTaskRunRow,
+    ScheduledTaskSubscriptionRow,
+    SchedulerThreadRow,
+)
 from deerflow.persistence.scheduled_tasks.sql import ScheduledTaskRepository

 __all__ = [
@@ -14,6 +20,7 @@ __all__ = [
    "ScheduledTaskRow",
    "ScheduledTaskRunRow",
    "ScheduledTaskStore",
+    "ScheduledTaskSubscriptionRow",
    "SchedulerThreadRow",
    "make_scheduled_task_store",
 ]
@@ -63,3 +63,46 @@ class ScheduledTaskStore(ABC):
    @abstractmethod
    async def update_delivery_profile(self, user_id: str, data: dict[str, Any]) -> dict[str, Any]:
        raise NotImplementedError
+
+    # --- Publish / subscribe ---
+
+    @abstractmethod
+    async def get_task_any(self, task_id: str) -> dict[str, Any] | None:
+        """Return a task by id without restricting by owner. Used for publish/subscribe lookups."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get_run_any(self, run_id: str) -> dict[str, Any] | None:
+        """Return a run by id without restricting by owner. Used for subscriber visibility."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def set_task_published(self, task_id: str, owner_user_id: str, *, published: bool) -> dict[str, Any] | None:
+        """Toggle published flag. Returns updated task or ``None`` if missing / not owned."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def list_published_tasks(self, *, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def list_subscribed_tasks(self, user_id: str) -> list[dict[str, Any]]:
+        """Return tasks the user has subscribed to. Each item carries ``_subscription`` dict."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get_subscription(self, task_id: str, user_id: str) -> dict[str, Any] | None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def upsert_subscription(self, task_id: str, user_id: str, *, notify_element: bool) -> dict[str, Any]:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete_subscription(self, task_id: str, user_id: str) -> bool:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def list_notify_subscribers(self, task_id: str) -> list[dict[str, Any]]:
+        """Subscribers with ``notify_element=True`` for the given task."""
+        raise NotImplementedError
@@ -2,7 +2,8 @@

 from __future__ import annotations

-from datetime import datetime
+import uuid
+from datetime import UTC, datetime
 from typing import Any

 from deerflow.persistence.scheduled_tasks.base import ScheduledTaskStore
@@ -15,6 +16,7 @@ class MemoryScheduledTaskStore(ScheduledTaskStore):
        self._runs_by_key: set[tuple[str, str]] = set()
        self._runs: dict[str, dict[str, Any]] = {}
        self._profiles: dict[str, dict[str, Any]] = {}
+        self._subscriptions: dict[tuple[str, str], dict[str, Any]] = {}

    async def get_or_create_scheduler_thread(self, user_id: str, thread_id_factory) -> dict:
        if user_id not in self._threads:
@@ -105,3 +107,73 @@ class MemoryScheduledTaskStore(ScheduledTaskStore):
        if "element_room_id" in data:
            profile["element_room_id"] = data["element_room_id"] or None
        return dict(profile)
+
+    # --- Publish / subscribe ---
+
+    async def get_task_any(self, task_id: str) -> dict[str, Any] | None:
+        task = self._tasks.get(task_id)
+        return dict(task) if task is not None else None
+
+    async def get_run_any(self, run_id: str) -> dict[str, Any] | None:
+        run = self._runs.get(run_id)
+        return dict(run) if run is not None else None
+
+    async def set_task_published(self, task_id: str, owner_user_id: str, *, published: bool) -> dict[str, Any] | None:
+        task = self._tasks.get(task_id)
+        if task is None or task["user_id"] != owner_user_id:
+            return None
+        task["published"] = bool(published)
+        task["published_at"] = datetime.now(UTC) if published else None
+        return dict(task)
+
+    async def list_published_tasks(self, *, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:
+        tasks = [dict(task) for task in self._tasks.values() if task.get("published")]
+        tasks.sort(key=lambda t: (t.get("published_at") or t.get("created_at") or ""), reverse=True)
+        return tasks[max(0, offset) : max(0, offset) + max(1, limit)]
+
+    async def list_subscribed_tasks(self, user_id: str) -> list[dict[str, Any]]:
+        out: list[dict[str, Any]] = []
+        for (task_id, sub_user_id), sub in self._subscriptions.items():
+            if sub_user_id != user_id:
+                continue
+            task = self._tasks.get(task_id)
+            if task is None:
+                continue
+            payload = dict(task)
+            payload["_subscription"] = dict(sub)
+            out.append(payload)
+        out.sort(key=lambda item: item["_subscription"].get("created_at") or "", reverse=True)
+        return out
+
+    async def get_subscription(self, task_id: str, user_id: str) -> dict[str, Any] | None:
+        sub = self._subscriptions.get((task_id, user_id))
+        return dict(sub) if sub is not None else None
+
+    async def upsert_subscription(self, task_id: str, user_id: str, *, notify_element: bool) -> dict[str, Any]:
+        key = (task_id, user_id)
+        existing = self._subscriptions.get(key)
+        now = datetime.now(UTC)
+        if existing is None:
+            sub = {
+                "id": str(uuid.uuid4()),
+                "task_id": task_id,
+                "user_id": user_id,
+                "notify_element": bool(notify_element),
+                "created_at": now,
+                "updated_at": now,
+            }
+            self._subscriptions[key] = sub
+            return dict(sub)
+        existing["notify_element"] = bool(notify_element)
+        existing["updated_at"] = now
+        return dict(existing)
+
+    async def delete_subscription(self, task_id: str, user_id: str) -> bool:
+        return self._subscriptions.pop((task_id, user_id), None) is not None
+
+    async def list_notify_subscribers(self, task_id: str) -> list[dict[str, Any]]:
+        return [
+            dict(sub)
+            for (tid, _uid), sub in self._subscriptions.items()
+            if tid == task_id and sub.get("notify_element")
+        ]
@@ -39,6 +39,20 @@ class ScheduledTaskRow(Base):
    last_run_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    last_status: Mapped[str | None] = mapped_column(String(32))
    execution_context_json: Mapped[dict] = mapped_column(JSON, default=dict)
+    published: Mapped[bool] = mapped_column(Boolean, default=False, index=True)
+    published_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
+    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC))
+
+
+class ScheduledTaskSubscriptionRow(Base):
+    __tablename__ = "scheduled_task_subscriptions"
+    __table_args__ = (UniqueConstraint("task_id", "user_id", name="uq_scheduled_task_subscriptions_task_user"),)
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(64), index=True)
+    user_id: Mapped[str] = mapped_column(String(64), index=True)
+    notify_element: Mapped[bool] = mapped_column(Boolean, default=False)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC))

@@ -5,17 +5,25 @@ from __future__ import annotations
 from datetime import UTC, datetime
 from typing import Any

-from sqlalchemy import select, update
+import uuid
+
+from sqlalchemy import delete, select, update
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

 from deerflow.persistence.scheduled_tasks.base import ScheduledTaskStore
-from deerflow.persistence.scheduled_tasks.model import ScheduledTaskDeliveryProfileRow, ScheduledTaskRow, ScheduledTaskRunRow, SchedulerThreadRow
+from deerflow.persistence.scheduled_tasks.model import (
+    ScheduledTaskDeliveryProfileRow,
+    ScheduledTaskRow,
+    ScheduledTaskRunRow,
+    ScheduledTaskSubscriptionRow,
+    SchedulerThreadRow,
+)


 def _to_dict(row) -> dict[str, Any]:
    data = row.to_dict()
-    for key in ("created_at", "updated_at", "next_run_at", "last_run_at", "scheduled_for", "started_at", "finished_at"):
+    for key in ("created_at", "updated_at", "next_run_at", "last_run_at", "scheduled_for", "started_at", "finished_at", "published_at"):
        value = data.get(key)
        if isinstance(value, datetime):
            data[key] = value.isoformat()
@@ -194,3 +202,107 @@ class ScheduledTaskRepository(ScheduledTaskStore):
            await session.commit()
            await session.refresh(row)
            return _to_dict(row)
+
+    # --- Publish / subscribe ---
+
+    async def get_task_any(self, task_id: str) -> dict[str, Any] | None:
+        async with self._sf() as session:
+            row = await session.get(ScheduledTaskRow, task_id)
+            return _to_dict(row) if row is not None else None
+
+    async def get_run_any(self, run_id: str) -> dict[str, Any] | None:
+        async with self._sf() as session:
+            row = await session.get(ScheduledTaskRunRow, run_id)
+            return _to_dict(row) if row is not None else None
+
+    async def set_task_published(self, task_id: str, owner_user_id: str, *, published: bool) -> dict[str, Any] | None:
+        async with self._sf() as session:
+            row = await session.get(ScheduledTaskRow, task_id)
+            if row is None or row.user_id != owner_user_id:
+                return None
+            row.published = bool(published)
+            row.published_at = datetime.now(UTC) if published else None
+            row.updated_at = datetime.now(UTC)
+            await session.commit()
+            await session.refresh(row)
+            return _to_dict(row)
+
+    async def list_published_tasks(self, *, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:
+        async with self._sf() as session:
+            result = await session.execute(
+                select(ScheduledTaskRow)
+                .where(ScheduledTaskRow.published.is_(True))
+                .order_by(ScheduledTaskRow.published_at.desc().nullslast(), ScheduledTaskRow.created_at.desc())
+                .offset(max(0, offset))
+                .limit(max(1, limit))
+            )
+            return [_to_dict(row) for row in result.scalars()]
+
+    async def list_subscribed_tasks(self, user_id: str) -> list[dict[str, Any]]:
+        async with self._sf() as session:
+            result = await session.execute(
+                select(ScheduledTaskRow, ScheduledTaskSubscriptionRow)
+                .join(ScheduledTaskSubscriptionRow, ScheduledTaskSubscriptionRow.task_id == ScheduledTaskRow.task_id)
+                .where(ScheduledTaskSubscriptionRow.user_id == user_id)
+                .order_by(ScheduledTaskSubscriptionRow.created_at.desc())
+            )
+            tasks: list[dict[str, Any]] = []
+            for task_row, sub_row in result.all():
+                task = _to_dict(task_row)
+                task["_subscription"] = _to_dict(sub_row)
+                tasks.append(task)
+            return tasks
+
+    async def get_subscription(self, task_id: str, user_id: str) -> dict[str, Any] | None:
+        async with self._sf() as session:
+            result = await session.execute(
+                select(ScheduledTaskSubscriptionRow)
+                .where(ScheduledTaskSubscriptionRow.task_id == task_id, ScheduledTaskSubscriptionRow.user_id == user_id)
+                .limit(1)
+            )
+            row = result.scalar_one_or_none()
+            return _to_dict(row) if row is not None else None
+
+    async def upsert_subscription(self, task_id: str, user_id: str, *, notify_element: bool) -> dict[str, Any]:
+        async with self._sf() as session:
+            result = await session.execute(
+                select(ScheduledTaskSubscriptionRow)
+                .where(ScheduledTaskSubscriptionRow.task_id == task_id, ScheduledTaskSubscriptionRow.user_id == user_id)
+                .limit(1)
+            )
+            row = result.scalar_one_or_none()
+            if row is None:
+                row = ScheduledTaskSubscriptionRow(
+                    id=str(uuid.uuid4()),
+                    task_id=task_id,
+                    user_id=user_id,
+                    notify_element=bool(notify_element),
+                )
+                session.add(row)
+            else:
+                row.notify_element = bool(notify_element)
+                row.updated_at = datetime.now(UTC)
+            await session.commit()
+            await session.refresh(row)
+            return _to_dict(row)
+
+    async def delete_subscription(self, task_id: str, user_id: str) -> bool:
+        async with self._sf() as session:
+            result = await session.execute(
+                delete(ScheduledTaskSubscriptionRow).where(
+                    ScheduledTaskSubscriptionRow.task_id == task_id,
+                    ScheduledTaskSubscriptionRow.user_id == user_id,
+                )
+            )
+            await session.commit()
+            return (result.rowcount or 0) > 0
+
+    async def list_notify_subscribers(self, task_id: str) -> list[dict[str, Any]]:
+        async with self._sf() as session:
+            result = await session.execute(
+                select(ScheduledTaskSubscriptionRow).where(
+                    ScheduledTaskSubscriptionRow.task_id == task_id,
+                    ScheduledTaskSubscriptionRow.notify_element.is_(True),
+                )
+            )
+            return [_to_dict(row) for row in result.scalars()]
@@ -197,16 +197,70 @@ class ScheduledTaskService:
        return thread

    async def list_tasks(self, user_id: str) -> list[dict[str, Any]]:
-        return await self.store.list_tasks(user_id)
+        owned = await self.store.list_tasks(user_id)
+        for task in owned:
+            task["role"] = "owner"
+        subscribed = await self.store.list_subscribed_tasks(user_id)
+        owned_ids = {task["task_id"] for task in owned}
+        for task in subscribed:
+            if task["task_id"] in owned_ids:
+                continue
+            task["role"] = "subscriber"
+            owned.append(task)
+        return owned
+
+    async def list_published(self, *, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:
+        """Public discovery list. Strips author prompt details to a short summary."""
+        tasks = await self.store.list_published_tasks(limit=limit, offset=offset)
+        result: list[dict[str, Any]] = []
+        for task in tasks:
+            prompt = (task.get("prompt") or "").strip()
+            summary = prompt if len(prompt) <= 200 else prompt[:200].rstrip() + "…"
+            result.append(
+                {
+                    "task_id": task["task_id"],
+                    "name": task.get("name"),
+                    "prompt_summary": summary,
+                    "schedule_text": task.get("schedule_text"),
+                    "timezone": task.get("timezone"),
+                    "published_at": task.get("published_at"),
+                    "owner_user_id": task.get("user_id"),
+                }
+            )
+        return result
+
+    async def _accessible_task(self, user_id: str, task_id: str) -> dict[str, Any] | None:
+        task = await self.store.get_task_any(task_id)
+        if task is None:
+            return None
+        if task["user_id"] == user_id:
+            task["role"] = "owner"
+            return task
+        sub = await self.store.get_subscription(task_id, user_id)
+        if sub is None:
+            return None
+        task["role"] = "subscriber"
+        task["_subscription"] = sub
+        return task

    async def list_runs(self, user_id: str, task_id: str, *, limit: int = 50) -> list[dict[str, Any]]:
-        task = await self.store.get_task(task_id, user_id)
-        if task is None:
+        accessible = await self._accessible_task(user_id, task_id)
+        if accessible is None:
            return []
-        return await self.store.list_runs(task_id, user_id, limit=limit)
+        return await self.store.list_runs(task_id, accessible["user_id"], limit=limit)

    async def get_run(self, user_id: str, run_id: str) -> dict[str, Any] | None:
-        return await self.store.get_run(run_id, user_id)
+        run = await self.store.get_run(run_id, user_id)
+        if run is not None:
+            return run
+        # Subscriber visibility: look up the run unfiltered, then verify the
+        # caller has a subscription on the parent task.
+        candidate = await self.store.get_run_any(run_id)
+        if candidate is None:
+            return None
+        if await self.store.get_subscription(candidate["task_id"], user_id) is None:
+            return None
+        return candidate

    async def delete_run(self, user_id: str, run_id: str) -> bool:
        deleted = await self.store.delete_run(run_id, user_id)
@@ -655,15 +709,62 @@ class ScheduledTaskService:
    async def _deliver_element_result(self, task: dict[str, Any], result_text: str) -> None:
        if self._element_client is None:
            return
-        profile = await self.store.get_delivery_profile(task["user_id"])
+        body = f"【定时任务】{task['name']}\n\n{result_text}"
+        # 1) Author
+        await self._deliver_to_user(task["user_id"], task, body)
+        # 2) Subscribers who opted into notifications
+        try:
+            subscribers = await self.store.list_notify_subscribers(task["task_id"])
+        except Exception:
+            logger.exception("Failed to list scheduled task subscribers for %s", task["task_id"])
+            subscribers = []
+        for sub in subscribers:
+            sub_user_id = sub.get("user_id")
+            if not sub_user_id or sub_user_id == task["user_id"]:
+                continue
+            await self._deliver_to_user(sub_user_id, task, body)
+
+    async def _deliver_to_user(self, user_id: str, task: dict[str, Any], body: str) -> None:
+        if self._element_client is None:
+            return
+        try:
+            profile = await self.store.get_delivery_profile(user_id)
+        except Exception:
+            logger.exception("Failed to load delivery profile for user %s", user_id)
+            return
        element_user_id = profile.get("element_user_id")
        if not isinstance(element_user_id, str) or not element_user_id.strip():
            return
        existing_room_id: str | None = profile.get("element_room_id") or None
-        body = f"【定时任务】{task['name']}\n\n{result_text}"
        try:
            room_id = await self._element_client.send_direct_message(element_user_id.strip(), body, room_id=existing_room_id)
            if room_id and room_id != existing_room_id:
-                await self.store.update_delivery_profile(task["user_id"], {"element_room_id": room_id})
+                await self.store.update_delivery_profile(user_id, {"element_room_id": room_id})
        except Exception:
            logger.exception("Failed to deliver scheduled task %s to Element user %s", task["task_id"], element_user_id)
+
+    # --- Publish / subscribe ---
+
+    async def publish_task(self, user_id: str, task_id: str) -> dict[str, Any] | None:
+        return await self.store.set_task_published(task_id, user_id, published=True)
+
+    async def unpublish_task(self, user_id: str, task_id: str) -> dict[str, Any] | None:
+        return await self.store.set_task_published(task_id, user_id, published=False)
+
+    async def subscribe_task(
+        self, user_id: str, task_id: str, *, notify_element: bool
+    ) -> dict[str, Any] | None:
+        task = await self.store.get_task_any(task_id)
+        if task is None:
+            return None
+        if task["user_id"] == user_id:
+            # Authors don't need to subscribe to their own tasks.
+            return None
+        existing = await self.store.get_subscription(task_id, user_id)
+        if existing is None and not task.get("published"):
+            # First-time subscription requires the task to currently be published.
+            return None
+        return await self.store.upsert_subscription(task_id, user_id, notify_element=notify_element)
+
+    async def unsubscribe_task(self, user_id: str, task_id: str) -> bool:
+        return await self.store.delete_subscription(task_id, user_id)