deerflow-code/offline-backend-20260512/backend/app/maf_roundtable/dispatcher.py

"""MAF run dispatcher — launches WorkflowRunner after create_run succeeds.

The dispatcher is the bridge between the HTTP layer (router) and the workflow
execution layer (WorkflowRunner). It:

1. Checks MAF SDK availability BEFORE creating the run (so the user gets a
   clean 503 rather than a run that will fail).
2. Launches ``WorkflowRunner.execute`` as a background asyncio task.
3. Tracks task handles for cancellation.
4. **Resume**: after an approval decision, the router calls
   ``dispatcher.resume(run_id, user_id)`` to restore the workflow from its
   checkpoint.
5. **Cancel cascade**: ``dispatcher.cancel(run_id)`` cancels both the
   background task AND the DeerFlow sub-runs. If the task handle is gone
   (process restart), it rebuilds a runner from persisted ``seat_runs`` and
   cascades cancel to DeerFlow.
6. Never silently degrades to sequential DeerFlow execution — if the SDK is
   missing or the workflow fails, the run is marked ``failed``.

Usage (from the router)::

    from app.maf_roundtable.dispatcher import get_dispatcher
    dispatcher = get_dispatcher(request.app)
    await dispatcher.dispatch(session, run, user_id=user_id)
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any

logger = logging.getLogger(__name__)


class MafSdkNotAvailableError(RuntimeError):
    """Raised when the MAF SDK is not installed and a run is requested."""


class MafRunDispatcher:
    """Manages background MAF workflow execution tasks."""

    def __init__(self, app: Any) -> None:
        self._app = app
        self._tasks: dict[str, asyncio.Task] = {}

    def check_sdk_available(self) -> None:
        """Raise MafSdkNotAvailableError if the MAF SDK is not installed."""
        try:
            import agent_framework  # noqa: F401
        except ImportError as exc:
            raise MafSdkNotAvailableError(
                "agent_framework SDK is not installed; run 'uv sync --extra maf' to enable MAF workflows"
            ) from exc

    async def dispatch(self, session: dict[str, Any], run: dict[str, Any], *, user_id: str) -> None:
        """Launch the workflow for a newly created run as a background task.

        The task is tracked so it can be cancelled via ``cancel(run_id)``.
        """
        from app.maf_roundtable.workflow_runner import WorkflowRunner

        run_id = run["id"]
        service = getattr(self._app.state, "maf_roundtable_service", None)
        if service is None:
            raise MafSdkNotAvailableError("MAF roundtable service not available")

        runner = WorkflowRunner(self._app, service=service, user_id=user_id)

        task = asyncio.create_task(self._execute(runner, session, run, run_id))
        self._tasks[run_id] = task
        task.add_done_callback(lambda t, rid=run_id: self._tasks.pop(rid, None))

    async def wait_for_run(self, run_id: str) -> None:
        """Wait for an in-process run without exposing its task to callers.

        Conversation V2 uses this only for a server-side status bridge.  The
        persisted run remains authoritative, so a missing task (for example
        after a worker restart) simply returns and lets recovery own the run.
        """
        task = self._tasks.get(run_id)
        if task is None:
            return
        try:
            await asyncio.shield(task)
        except asyncio.CancelledError:
            raise
        except Exception:  # noqa: BLE001 - runner already persisted failure
            return

    async def _execute(self, runner: Any, session: dict[str, Any], run: dict[str, Any], run_id: str) -> None:
        """Execute the workflow, catching all exceptions (they're already persisted by the runner)."""
        try:
            await runner.execute(session, run)
        except Exception:
            # The runner already persists the failure. Log for observability.
            logger.exception("MAF workflow execution failed for run %s", run_id)

    async def resume(self, run_id: str, *, user_id: str) -> dict[str, Any] | None:
        """Resume a workflow from its checkpoint after an approval decision.

        Called by the router after ``decide_approval`` succeeds. Loads the
        runner (rebuilding if the process restarted), then calls
        ``runner.resume(run_id)`` which restores ``RoundtableState.from_dict()``
        and executes the approved action (at-most-once).
        """
        from app.maf_roundtable.workflow_runner import WorkflowRunner

        service = getattr(self._app.state, "maf_roundtable_service", None)
        if service is None:
            raise MafSdkNotAvailableError("MAF roundtable service not available")

        runner = WorkflowRunner(self._app, service=service, user_id=user_id)
        return await runner.resume(run_id)

    async def recover(self, run_id: str, *, user_id: str) -> dict[str, Any] | None:
        """Re-enter the workflow from checkpoint after a process restart.

        Unlike ``resume`` (approval path), this continues orchestration for
        ``running`` runs and skips ``completed_nodes``.
        """
        from app.maf_roundtable.workflow_runner import WorkflowRunner

        service = getattr(self._app.state, "maf_roundtable_service", None)
        if service is None:
            raise MafSdkNotAvailableError("MAF roundtable service not available")

        # Avoid overlapping with an in-process execute task.
        existing = self._tasks.get(run_id)
        if existing is not None and not existing.done():
            logger.info("recover(%s): background task still running — skip", run_id)
            return await service.store.runs.get_run(run_id)

        runner = WorkflowRunner(self._app, service=service, user_id=user_id)
        task = asyncio.create_task(self._recover(runner, run_id))
        self._tasks[run_id] = task
        task.add_done_callback(lambda t, rid=run_id: self._tasks.pop(rid, None))
        return await service.store.runs.get_run(run_id)

    async def _recover(self, runner: Any, run_id: str) -> None:
        try:
            await runner.recover(run_id)
        except Exception:
            logger.exception("MAF recover failed for run %s", run_id)

    async def cancel(self, run_id: str) -> dict[str, Any]:
        """Cancel a running workflow AND cascade to DeerFlow sub-runs.

        Returns a structured cancel result so the router can decide whether to
        finalize the parent run as ``cancelled`` (all sub-runs confirmed
        cancelled / already terminal) or leave it non-terminal (some sub-runs
        failed to cancel — must not masquerade as cancelled).

        Result shape::

            {
                "all_cancelled": bool,
                "cancelled_run_ids": [...],
                "failed_run_ids": [...],
                "still_active_run_ids": [...],
            }
        """
        from app.maf_roundtable.workflow_runner import WorkflowRunner

        _EMPTY = {"all_cancelled": False, "cancelled_run_ids": [], "failed_run_ids": [], "still_active_run_ids": []}

        # P1c cancel order: the parent run is already in ``cancelling`` (the
        # router sets it before calling here). We then (1) cascade-cancel the
        # DeerFlow seat runs first and confirm/record the outcome, and only
        # (2) cancel the local MAF asyncio task LAST. The router finalizes the
        # parent run as ``cancelled`` only when every sub-run confirmed
        # cancelled — a failed cascade must never masquerade as cancelled.
        service = getattr(self._app.state, "maf_roundtable_service", None)
        if service is None:
            return _EMPTY

        run = await service.store.runs.get_run(run_id)
        if run is None:
            return _EMPTY

        session = await service.store.sessions.get_session(run["session_id"], user_id=None)
        if session is None:
            return _EMPTY

        user_id = session.get("user_id", "default")
        runner = WorkflowRunner(self._app, service=service, user_id=user_id)

        # Step 1: cascade cancel to DeerFlow sub-runs (works even after a
        # process restart — the runner reads persisted seat_runs).
        try:
            result = await runner.cancel(run_id)
        except Exception:  # noqa: BLE001
            logger.exception("MAF cancel cascade failed for run %s", run_id)
            result = dict(_EMPTY)

        # Step 2: cancel the local MAF asyncio task LAST (after the external
        # sub-runs have been addressed), so we never drop the orchestrator
        # before its sub-runs are cancelled/recorded.
        task = self._tasks.get(run_id)
        if task is not None and not task.done():
            task.cancel()
            try:
                await task
            except (asyncio.CancelledError, Exception):  # noqa: BLE001
                pass

        return result


_dispatchers: dict[int, MafRunDispatcher] = {}


def get_dispatcher(app: Any) -> MafRunDispatcher:
    """Get or create the dispatcher for an app instance (keyed by id)."""
    app_id = id(app)
    if app_id not in _dispatchers:
        _dispatchers[app_id] = MafRunDispatcher(app)
    return _dispatchers[app_id]