"""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]