dify/api/services/workflow_collaboration_service.py

from __future__ import annotations

import logging
import os
import socket
import time
import uuid
from collections.abc import Mapping
from typing import Any, NamedTuple, override

from socketio.exceptions import TimeoutError as SocketIOTimeoutError  # type: ignore[reportMissingTypeStubs]
from sqlalchemy import select
from sqlalchemy.orm import Session

from configs import dify_config
from core.rbac import RBACPermission, RBACResourceScope
from models.account import Account
from models.model import App
from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo
from services.enterprise.rbac_service import RBACService
from services.feature_service import FeatureService

logger = logging.getLogger(__name__)

SERVER_HEARTBEAT_INTERVAL_SECONDS = 30
SYNC_REQUEST_TIMEOUT_SECONDS = 15
_PROCESS_SERVER_ID: str | None = None
_PROCESS_SERVER_ID_PID: int | None = None


class WorkflowCollaborationNotEntitledError(Exception):
    """The workspace cannot use workflow collaboration."""


class WorkflowPageResume(NamedTuple):
    """Validated page credential and join sequence supplied by the browser."""

    page_id: str
    sequence: int


def _get_process_server_id() -> str:
    global _PROCESS_SERVER_ID, _PROCESS_SERVER_ID_PID  # pylint: disable=global-statement

    pid = os.getpid()
    if _PROCESS_SERVER_ID is None or pid != _PROCESS_SERVER_ID_PID:
        _PROCESS_SERVER_ID_PID = pid
        _PROCESS_SERVER_ID = f"{socket.gethostname()}:{pid}:{uuid.uuid4().hex}"
    return _PROCESS_SERVER_ID


class WorkflowCollaborationService:
    """
    Coordinate workflow collaboration state across Socket.IO workers.

    Socket.IO rooms are process-local unless backed by a message queue, while online users and leader election live in
    Redis. Each websocket worker writes a small heartbeat keyed by `server_id`; session rows store their owner so a
    worker can distinguish a live remote sid from a stale sid left behind by a dead worker. Visibility events are
    ordered per socket, and sync requests wait for the selected saver to acknowledge its draft write.
    """

    _heartbeat_started: bool
    _repository: WorkflowCollaborationRepository
    _server_id_override: str | None
    _socketio: Any

    def __init__(
        self, repository: WorkflowCollaborationRepository, socketio: Any, server_id: str | None = None
    ) -> None:
        self._repository = repository
        self._socketio = socketio
        self._server_id_override = server_id
        self._heartbeat_started = False

    @override
    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(repository={self._repository})"

    @property
    def server_id(self) -> str:
        return self._server_id_override or _get_process_server_id()

    def _refresh_server_heartbeat(self) -> None:
        self._repository.refresh_server_heartbeat(self.server_id)

    def _server_heartbeat_loop(self) -> None:
        while True:
            try:
                self._refresh_server_heartbeat()
                self._repository.refresh_server_sessions(self.server_id)
            except Exception:
                logger.exception("Failed to refresh workflow collaboration server heartbeat")
            self._socketio.sleep(SERVER_HEARTBEAT_INTERVAL_SECONDS)

    def _ensure_server_heartbeat_started(self) -> None:
        if self._heartbeat_started:
            return

        self._heartbeat_started = True
        self._refresh_server_heartbeat()
        self._socketio.start_background_task(self._server_heartbeat_loop)

    def save_socket_identity(self, sid: str, user: Account) -> None:
        """Persist the authenticated console user on the raw socket session."""
        self._socketio.save_session(
            sid,
            {
                "user_id": user.id,
                "username": user.name,
                "avatar": user.avatar,
                "tenant_id": user.current_tenant_id,
            },
        )

    def authorize_and_join_workflow_room(
        self, workflow_id: str, sid: str, *, session: Session, resume: WorkflowPageResume | None = None
    ) -> tuple[str, bool] | None:
        """
        Join a collaboration room only after validating the socket session, app access, and workspace entitlement.

        The Socket.IO payload still calls the room key `workflow_id`, but the identifier is the workflow app's
        `App.id`. Missing access or an outdated join returns `None`; an ineligible workspace raises
        `WorkflowCollaborationNotEntitledError`. An admitted resume supersedes only the prior SID for
        this page in the authenticated workflow and account.
        """
        socket_session = self._socketio.get_session(sid)
        user_id = socket_session.get("user_id")
        tenant_id = socket_session.get("tenant_id")
        if not user_id or not tenant_id:
            return None

        if not self._can_access_workflow(workflow_id, str(tenant_id), str(user_id), session=session):
            logger.warning(
                "Workflow collaboration join rejected: workflow_id=%s tenant_id=%s user_id=%s sid=%s",
                workflow_id,
                tenant_id,
                user_id,
                sid,
            )
            return None

        if not FeatureService.can_use_workflow_collaboration(str(tenant_id)):
            raise WorkflowCollaborationNotEntitledError

        self._ensure_server_heartbeat_started()

        session_info: WorkflowSessionInfo = {
            "user_id": str(user_id),
            "username": str(socket_session.get("username", "Unknown")),
            "avatar": socket_session.get("avatar"),
            "sid": sid,
            "connected_at": int(time.time()),
            "server_id": self.server_id,
            # Joins are assumed visible; hidden tabs re-report via a graph_view_state
            # event right after receiving the post-join "status" emit.
            "graph_active": True,
        }

        previous_sid = None
        replaced_sid = None
        with self._repository.room_state_lock(workflow_id):
            if resume is not None:
                previous = self._repository.get_page_session(workflow_id, session_info["user_id"], resume.page_id)
                if previous is not None:
                    if resume.sequence < previous["sequence"]:
                        return None
                    if previous["sid"] != sid:
                        if resume.sequence == previous["sequence"]:
                            return None
                        previous_sid = previous["sid"]
                        mapping = self._repository.get_sid_mapping(previous_sid)
                        if mapping is not None and (
                            mapping["workflow_id"],
                            mapping["user_id"],
                            mapping.get("page_id"),
                        ) == (workflow_id, session_info["user_id"], resume.page_id):
                            replaced_sid = previous_sid

            self._socketio.enter_room(sid, workflow_id)
            self._repository.set_session_info(workflow_id, session_info, page_id=resume.page_id if resume else None)
            if resume is not None:
                self._repository.set_page_session(
                    workflow_id, session_info["user_id"], resume.page_id, {"sid": sid, "sequence": resume.sequence}
                )
            if replaced_sid is not None:
                self._repository.delete_session(workflow_id, replaced_sid)
            if previous_sid is not None:
                self.handle_leader_disconnect(workflow_id, previous_sid, joining_sid=sid)

            self.get_or_set_leader(workflow_id, sid)
            is_leader = self.broadcast_online_users(workflow_id, joining_sid=sid) == sid
            status = {"isLeader": is_leader}
            if is_leader and replaced_sid is not None:
                status["requiresGraphReload"] = True
            self._socketio.emit("status", status, room=sid)

        # Socket.IO can run the disconnect handler synchronously. Room state already fences this SID.
        if replaced_sid is not None:
            try:
                self._socketio.disconnect(replaced_sid)
            except Exception:
                logger.exception("Failed to close superseded workflow collaboration session: sid=%s", replaced_sid)

        return str(user_id), is_leader

    def _can_access_workflow(self, workflow_id: str, tenant_id: str, user_id: str, *, session: Session) -> bool:
        """Check tenant and app permission without relying on Flask's scoped session."""
        with session.begin():
            app = session.execute(
                select(App.id, App.maintainer).where(
                    App.id == workflow_id, App.tenant_id == tenant_id, App.status == "normal"
                )
            ).one_or_none()
        if app is None:
            return False

        app_id, maintainer = app
        if not dify_config.RBAC_ENABLED or maintainer == user_id:
            return True
        return RBACService.CheckAccess.check(
            tenant_id,
            user_id,
            scene=RBACPermission.APP_EDIT,
            resource_type=RBACResourceScope.APP,
            resource_id=app_id,
        )

    def disconnect_session(self, sid: str) -> str | None:
        """Remove a known session and return its workflow ID for final draft cleanup.

        Unknown and superseded sessions return None. Cleanup is serialized with joins,
        so a late disconnect cannot act on the replacement session's room state.
        """
        mapping = self._repository.get_sid_mapping(sid)
        if not mapping:
            return None

        workflow_id = mapping["workflow_id"]
        with self._repository.room_state_lock(workflow_id):
            if self._repository.get_sid_mapping(sid) != mapping:
                return None
            self._repository.delete_session(workflow_id, sid)
            self.handle_leader_disconnect(workflow_id, sid)
            self.broadcast_online_users(workflow_id)
        return workflow_id

    def relay_collaboration_event(self, sid: str, data: Mapping[str, object]) -> tuple[dict[str, object], int]:
        """Route collaboration control events, directing save and graph-resync requests to the active leader."""
        mapping = self._repository.get_sid_mapping(sid)
        if not mapping:
            return {"msg": "unauthorized"}, 401

        workflow_id = mapping["workflow_id"]
        user_id = mapping["user_id"]
        self.refresh_session_state(workflow_id, sid)

        event_type = data.get("type")
        event_data = data.get("data")
        timestamp = data.get("timestamp", int(time.time()))

        if not event_type:
            return {"msg": "invalid event type"}, 400

        if event_type == "graph_view_state":
            if not isinstance(event_data, Mapping):
                return {"msg": "invalid graph_view_state"}, 400

            graph_active = event_data.get("graphActive")
            sequence = event_data.get("sequence")
            if not isinstance(graph_active, bool):
                return {"msg": "invalid graph_view_state"}, 400
            if not isinstance(sequence, int) or isinstance(sequence, bool) or sequence < 0:
                return {"msg": "invalid graph_view_state"}, 400

            # Sequence the visibility write together with its leader side effect. Otherwise a newer
            # visible event can land after the Lua write but before an older hidden handler demotes.
            with self._repository.room_state_lock(workflow_id):
                applied = self._repository.update_session_graph_active(workflow_id, sid, graph_active, sequence)
                if not applied:
                    return {"msg": "graph_view_state_ignored"}, 200

                # Write the flag before re-electing so tier-1 selection already excludes the
                # session that just went hidden (including one _ensure_leader just promoted).
                if not graph_active:
                    self._demote_leader_if_hidden(workflow_id, sid)
            return {"msg": "graph_view_state_updated"}, 200

        if event_type == "sync_request":
            if not isinstance(event_data, Mapping):
                return {"msg": "invalid sync_request"}, 400
            request_id = event_data.get("requestId")
            if not isinstance(request_id, str) or not request_id.strip():
                return {"msg": "invalid sync_request"}, 400

            with self._repository.room_state_lock(workflow_id):
                if self._repository.get_sid_mapping(sid) != mapping:
                    return {"msg": "unauthorized"}, 401
                target_sid = self._select_sync_target(workflow_id, sid, require_visible_leader=True)

            if not target_sid:
                return {"msg": "no_active_leader", "requestId": request_id}, 503

            target_data = dict(event_data)
            target_data["requestId"] = request_id
            try:
                result = self._socketio.call(
                    "collaboration_update",
                    {"type": event_type, "userId": user_id, "data": target_data, "timestamp": timestamp},
                    to=target_sid,
                    timeout=SYNC_REQUEST_TIMEOUT_SECONDS,
                )
            except SocketIOTimeoutError:
                logger.warning(
                    "Workflow collaboration sync request timed out: workflow_id=%s requester_sid=%s target_sid=%s",
                    workflow_id,
                    sid,
                    target_sid,
                )
                return {"msg": "sync_request_timeout", "requestId": request_id}, 504

            if not isinstance(result, Mapping) or result.get("success") is not True:
                response: dict[str, object] = {
                    "msg": "workflow_sync_failed",
                    "requestId": request_id,
                    "success": False,
                }
                if isinstance(result, Mapping) and isinstance(result.get("error"), str):
                    response["error"] = result["error"]
                return response, 502

            workflow_hash = result.get("hash")
            updated_at = result.get("updatedAt")
            if not isinstance(workflow_hash, str) or not workflow_hash:
                return {"msg": "invalid_sync_response", "requestId": request_id}, 502
            if not isinstance(updated_at, int) or isinstance(updated_at, bool):
                return {"msg": "invalid_sync_response", "requestId": request_id}, 502

            return {
                "msg": "workflow_synced",
                "requestId": request_id,
                "success": True,
                "hash": workflow_hash,
                "updatedAt": updated_at,
            }, 200

        if event_type == "graph_resync_request":
            with self._repository.room_state_lock(workflow_id):
                if self._repository.get_sid_mapping(sid) != mapping:
                    return {"msg": "unauthorized"}, 401
                resync_target_sid = self._select_sync_target(workflow_id, sid, require_visible_leader=False)

            if not resync_target_sid:
                return {"msg": "no_active_leader"}, 503

            self._socketio.emit(
                "collaboration_update",
                {"type": event_type, "userId": user_id, "data": event_data, "timestamp": timestamp},
                to=resync_target_sid,
            )
            return {"msg": "graph_resync_request_forwarded"}, 200

        self._socketio.emit(
            "collaboration_update",
            {"type": event_type, "userId": user_id, "data": event_data, "timestamp": timestamp},
            room=workflow_id,
            skip_sid=sid,
        )

        return {"msg": "event_broadcasted"}, 200

    def _select_sync_target(self, workflow_id: str, sid: str, *, require_visible_leader: bool) -> str | None:
        """Route an admitted request under the room lock; publish any necessary leader change."""
        leader_sid = self._repository.get_current_leader(workflow_id)
        if leader_sid and self.is_session_active(workflow_id, leader_sid):
            if not require_visible_leader or self._is_session_graph_active(workflow_id, leader_sid):
                return leader_sid
            # A hidden tab's canvas may be stale even when its CRDT can serve snapshots.
            # Saves prefer a visible editor, or the requester if all tabs are hidden.
            target_sid = self._select_graph_leader(workflow_id, preferred_sid=sid) or leader_sid
        else:
            target_sid = self._select_graph_leader(workflow_id, preferred_sid=sid)

        if target_sid != leader_sid:
            if target_sid is None:
                self._repository.delete_leader(workflow_id)
            else:
                self._repository.set_leader(workflow_id, target_sid)
                self.broadcast_leader_change(workflow_id, target_sid)
        return target_sid

    def relay_graph_event(self, sid: str, data: object) -> tuple[dict[str, str], int]:
        mapping = self._repository.get_sid_mapping(sid)
        if not mapping:
            return {"msg": "unauthorized"}, 401

        workflow_id = mapping["workflow_id"]
        self.refresh_session_state(workflow_id, sid)

        self._socketio.emit("graph_update", data, room=workflow_id, skip_sid=sid)

        return {"msg": "graph_update_broadcasted"}, 200

    def get_or_set_leader(self, workflow_id: str, sid: str) -> str:
        """Keep the current leader or elect the joining SID while the room is locked."""
        current_leader = self._repository.get_current_leader(workflow_id)

        if current_leader:
            if self.is_session_active(workflow_id, current_leader):
                return current_leader
            self._repository.delete_session(workflow_id, current_leader)
            self._repository.delete_leader(workflow_id)

        was_set = self._repository.set_leader_if_absent(workflow_id, sid)

        if was_set:
            if current_leader:
                self.broadcast_leader_change(workflow_id, sid, skip_sid=sid)
            return sid

        current_leader = self._repository.get_current_leader(workflow_id)
        if current_leader:
            return current_leader

        return sid

    def handle_leader_disconnect(
        self, workflow_id: str, disconnected_sid: str, *, joining_sid: str | None = None
    ) -> None:
        """Transfer leadership to a remaining editor, leaving a new join to recover separately."""
        current_leader = self._repository.get_current_leader(workflow_id)
        if not current_leader:
            return

        if current_leader != disconnected_sid:
            return

        new_leader_sid = self._select_graph_leader(workflow_id, excluded_sid=joining_sid)
        if new_leader_sid:
            self._repository.set_leader(workflow_id, new_leader_sid)
            self.broadcast_leader_change(workflow_id, new_leader_sid, skip_sid=joining_sid)
        else:
            self._repository.delete_leader(workflow_id)

    def broadcast_leader_change(
        self, workflow_id: str, new_leader_sid: str | None, *, skip_sid: str | None = None
    ) -> None:
        """Notify existing editors. A joining SID receives its own status with recovery instructions."""
        for sid in self._repository.get_session_sids(workflow_id):
            if sid != skip_sid:
                try:
                    is_leader = new_leader_sid is not None and sid == new_leader_sid
                    self._socketio.emit("status", {"isLeader": is_leader}, room=sid)
                except Exception:
                    logging.exception("Failed to emit leader status to session %s", sid)

    def get_current_leader(self, workflow_id: str) -> str | None:
        return self._repository.get_current_leader(workflow_id)

    def _prune_inactive_sessions(self, workflow_id: str) -> list[WorkflowSessionInfo]:
        """Remove inactive sessions from storage and return active sessions only."""
        sessions = self._repository.list_sessions(workflow_id)
        if not sessions:
            return []

        active_sessions: list[WorkflowSessionInfo] = []
        stale_sids: list[str] = []
        for session in sessions:
            sid = session["sid"]
            if self.is_session_active(workflow_id, sid):
                active_sessions.append(session)
            else:
                stale_sids.append(sid)

        for sid in stale_sids:
            self._repository.delete_session(workflow_id, sid)

        return active_sessions

    def broadcast_online_users(self, workflow_id: str, *, joining_sid: str | None = None) -> str | None:
        """Prune and publish the online list under the room lock, returning its final leader."""
        users = self._prune_inactive_sessions(workflow_id)
        users.sort(key=lambda x: x.get("connected_at") or 0)

        leader_sid = self.get_current_leader(workflow_id)
        previous_leader = leader_sid
        active_sids = {user["sid"] for user in users}
        if leader_sid and leader_sid not in active_sids:
            self._repository.delete_leader(workflow_id)
            leader_sid = None

        if not leader_sid and users:
            leader_sid = self._select_graph_leader(workflow_id)
            if leader_sid:
                self._repository.set_leader(workflow_id, leader_sid)

        if leader_sid != previous_leader:
            self.broadcast_leader_change(workflow_id, leader_sid, skip_sid=joining_sid)

        self._socketio.emit(
            "online_users",
            {"workflow_id": workflow_id, "users": users, "leader": leader_sid},
            room=workflow_id,
        )
        return leader_sid

    def refresh_session_state(self, workflow_id: str, sid: str) -> None:
        self._refresh_server_heartbeat()
        self._repository.refresh_session_state(workflow_id, sid)
        self._ensure_leader(workflow_id, sid)

    def _ensure_leader(self, workflow_id: str, sid: str) -> None:
        """Keep a live leader; elect a requester only while its SID is still admitted."""
        current_leader = self._repository.get_current_leader(workflow_id)
        if current_leader and self.is_session_active(workflow_id, current_leader):
            self._repository.expire_leader(workflow_id)
            return

        with self._repository.room_state_lock(workflow_id):
            mapping = self._repository.get_sid_mapping(sid)
            if mapping is None or mapping["workflow_id"] != workflow_id:
                return
            current_leader = self._repository.get_current_leader(workflow_id)
            if current_leader and self.is_session_active(workflow_id, current_leader):
                self._repository.expire_leader(workflow_id)
                return

            self._repository.set_leader(workflow_id, sid)
            self.broadcast_leader_change(workflow_id, sid)

    def _select_graph_leader(
        self,
        workflow_id: str,
        preferred_sid: str | None = None,
        *,
        require_graph_active: bool = False,
        excluded_sid: str | None = None,
    ) -> str | None:
        """Pick a leader, preferring sessions whose canvas tab is visible.

        Hidden tabs freeze rAF-driven CRDT->canvas application, so a visible session is
        always the freshest saver. When every tab is hidden the room must still keep a
        leader (followers drop sync_requests unless they believe they are the leader),
        so tier 2 falls back to any active session unless require_graph_active is set.
        """
        active_sessions = [
            session
            for session in self._repository.list_sessions(workflow_id)
            if session["sid"] != excluded_sid and self.is_session_active(workflow_id, session["sid"])
        ]
        visible_sids = [session["sid"] for session in active_sessions if session.get("graph_active", True)]

        candidate_sids = visible_sids
        if not candidate_sids and not require_graph_active:
            candidate_sids = [session["sid"] for session in active_sessions]

        if not candidate_sids:
            return None
        if preferred_sid and preferred_sid in candidate_sids:
            return preferred_sid
        return candidate_sids[0]

    def _is_session_graph_active(self, workflow_id: str, sid: str) -> bool:
        """Default to True on missing/unreadable session data so read failures never churn leadership."""
        session_info = self._repository.get_session_info(workflow_id, sid)
        if session_info is None:
            return True
        return bool(session_info.get("graph_active", True))

    def _demote_leader_if_hidden(self, workflow_id: str, hidden_sid: str) -> None:
        current_leader = self._repository.get_current_leader(workflow_id)
        if current_leader != hidden_sid:
            return

        new_leader = self._select_graph_leader(workflow_id, require_graph_active=True)
        # No visible session: keep the hidden leader rather than leaving the room leaderless.
        if not new_leader or new_leader == hidden_sid:
            return

        self._repository.set_leader(workflow_id, new_leader)
        self.broadcast_leader_change(workflow_id, new_leader)

    def is_session_active(self, workflow_id: str, sid: str) -> bool:
        if not sid:
            return False

        mapping = self._repository.get_sid_mapping(sid)
        if not mapping or mapping["workflow_id"] != workflow_id:
            return False

        if not self._repository.session_exists(workflow_id, sid):
            return False

        server_id = mapping.get("server_id")
        if not server_id:
            return self._is_socket_connected_locally(sid)

        if server_id == self.server_id:
            return self._is_socket_connected_locally(sid)

        return self._repository.server_heartbeat_exists(server_id)

    def _is_socket_connected_locally(self, sid: str) -> bool:
        try:
            return bool(self._socketio.manager.is_connected(sid, "/"))
        except AttributeError:
            return False