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, override

from sqlalchemy import select
from sqlalchemy.orm import Session

from models.account import Account
from models.model import App
from repositories.workflow_collaboration_repository import WorkflowCollaborationRepository, WorkflowSessionInfo

logger = logging.getLogger(__name__)

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


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

    _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
    ) -> tuple[str, bool] | None:
        """
        Join a collaboration room only after validating the socket session and tenant-scoped app access.

        The Socket.IO payload still calls the room key `workflow_id`, but the identifier is the workflow app's
        `App.id`. Returning `None` lets the controller reject the join before any Redis or room state is created.
        """
        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), 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

        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,
        }

        self._repository.set_session_info(workflow_id, session_info)

        leader_sid = self.get_or_set_leader(workflow_id, sid)
        is_leader = leader_sid == sid

        self._socketio.enter_room(sid, workflow_id)
        self.broadcast_online_users(workflow_id)

        self._socketio.emit("status", {"isLeader": is_leader}, room=sid)

        return str(user_id), is_leader

    def _can_access_workflow(self, workflow_id: str, tenant_id: str, *, session: Session) -> bool:
        """Check room access without relying on Flask's app-context-bound scoped session."""
        app_id = session.scalar(select(App.id).where(App.id == workflow_id, App.tenant_id == tenant_id).limit(1))
        return app_id is not None

    def disconnect_session(self, sid: str) -> None:
        mapping = self._repository.get_sid_mapping(sid)
        if not mapping:
            return

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

        self.handle_leader_disconnect(workflow_id, sid)
        self.broadcast_online_users(workflow_id)

    def relay_collaboration_event(self, sid: str, data: Mapping[str, 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"]
        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 == "sync_request":
            leader_sid = self._repository.get_current_leader(workflow_id)
            target_sid: str | None
            if leader_sid and self.is_session_active(workflow_id, leader_sid):
                target_sid = leader_sid
            else:
                if leader_sid:
                    self._repository.delete_leader(workflow_id)
                target_sid = self._select_graph_leader(workflow_id, preferred_sid=sid)
                if target_sid:
                    self._repository.set_leader(workflow_id, target_sid)
                    self.broadcast_leader_change(workflow_id, target_sid)

            if not target_sid:
                return {"msg": "no_active_leader"}, 200

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

            return {"msg": "sync_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 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:
        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)
            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) -> None:
        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)
        if new_leader_sid:
            self._repository.set_leader(workflow_id, new_leader_sid)
            self.broadcast_leader_change(workflow_id, new_leader_sid)
        else:
            self._repository.delete_leader(workflow_id)

    def broadcast_leader_change(self, workflow_id: str, new_leader_sid: str | None) -> None:
        for sid in self._repository.get_session_sids(workflow_id):
            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) -> None:
        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)

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

    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:
        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

        if current_leader:
            self._repository.delete_leader(workflow_id)

        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) -> str | None:
        session_sids = [
            session["sid"]
            for session in self._repository.list_sessions(workflow_id)
            if session.get("graph_active", True) and self.is_session_active(workflow_id, session["sid"])
        ]
        if not session_sids:
            return None
        if preferred_sid and preferred_sid in session_sids:
            return preferred_sid
        return session_sids[0]

    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:
            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