dify/api/services/workflow_collaboration_service.py

from __future__ import annotations

import logging
import time
from collections.abc import Mapping

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


class WorkflowCollaborationService:
    def __init__(self, repository: WorkflowCollaborationRepository, socketio) -> None:
        self._repository = repository
        self._socketio = socketio

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

    def save_session(self, sid: str, user: Account) -> None:
        self._socketio.save_session(
            sid,
            {
                "user_id": user.id,
                "username": user.name,
                "avatar": user.avatar,
            },
        )

    def register_session(self, workflow_id: str, sid: str) -> tuple[str, bool] | None:
        session = self._socketio.get_session(sid)
        user_id = session.get("user_id")
        if not user_id:
            return None

        session_info: WorkflowSessionInfo = {
            "user_id": str(user_id),
            "username": str(session.get("username", "Unknown")),
            "avatar": session.get("avatar"),
            "sid": sid,
            "connected_at": int(time.time()),
            "graph_active": True,
            "active_skill_file_id": None,
        }

        self._repository.set_session_info(workflow_id, session_info)

        leader_sid = self.get_or_set_leader(workflow_id, sid)
        is_leader = leader_sid == sid if leader_sid else False

        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 disconnect_session(self, sid: str) -> None:
        mapping = self._repository.get_sid_mapping(sid)
        if not mapping:
            return

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

        self.handle_leader_disconnect(workflow_id, sid)
        if active_skill_file_id:
            self.handle_skill_leader_disconnect(workflow_id, active_skill_file_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 == "graph_view_active":
            is_active = False
            if isinstance(event_data, dict):
                is_active = bool(event_data.get("active") or False)
            self._repository.set_graph_active(workflow_id, sid, is_active)
            self.refresh_session_state(workflow_id, sid)
            self.broadcast_online_users(workflow_id)
            return {"msg": "graph_view_active_updated"}, 200

        if event_type == "skill_file_active":
            file_id = None
            is_active = False
            if isinstance(event_data, dict):
                file_id = event_data.get("file_id")
                is_active = bool(event_data.get("active") or False)

            if not file_id or not isinstance(file_id, str):
                return {"msg": "invalid skill_file_active payload"}, 400

            previous_file_id = self._repository.get_active_skill_file_id(workflow_id, sid)
            next_file_id = file_id if is_active else None

            if previous_file_id == next_file_id:
                self.refresh_session_state(workflow_id, sid)
                return {"msg": "skill_file_active_unchanged"}, 200

            self._repository.set_active_skill_file(workflow_id, sid, next_file_id)
            self.refresh_session_state(workflow_id, sid)

            if previous_file_id:
                self._ensure_skill_leader(workflow_id, previous_file_id)
            if next_file_id:
                self._ensure_skill_leader(workflow_id, next_file_id, preferred_sid=sid)

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

        if event_type == "sync_request":
            leader_sid = self._repository.get_current_leader(workflow_id)
            if leader_sid and (
                self.is_session_active(workflow_id, leader_sid)
                and self._repository.is_graph_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 relay_skill_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("skill_update", data, room=workflow_id, skip_sid=sid)

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

    def get_or_set_leader(self, workflow_id: str, sid: str) -> str | None:
        current_leader = self._repository.get_current_leader(workflow_id)

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

        new_leader_sid = self._select_graph_leader(workflow_id, preferred_sid=sid)
        if not new_leader_sid:
            return None

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

        if was_set:
            if current_leader:
                self.broadcast_leader_change(workflow_id, new_leader_sid)
            return new_leader_sid

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

        return new_leader_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)
            self.broadcast_leader_change(workflow_id, None)

    def handle_skill_leader_disconnect(self, workflow_id: str, file_id: str, disconnected_sid: str) -> None:
        current_leader = self._repository.get_skill_leader(workflow_id, file_id)
        if not current_leader:
            return

        if current_leader != disconnected_sid:
            return

        new_leader_sid = self._select_skill_leader(workflow_id, file_id)
        if new_leader_sid:
            self._repository.set_skill_leader(workflow_id, file_id, new_leader_sid)
            self.broadcast_skill_leader_change(workflow_id, file_id, new_leader_sid)
        else:
            self._repository.delete_skill_leader(workflow_id, file_id)
            self.broadcast_skill_leader_change(workflow_id, file_id, None)

    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 broadcast_skill_leader_change(self, workflow_id: str, file_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("skill_status", {"file_id": file_id, "isLeader": is_leader}, room=sid)
            except Exception:
                logging.exception("Failed to emit skill 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._repository.refresh_session_state(workflow_id, sid)
        self._ensure_leader(workflow_id, sid)
        active_skill_file_id = self._repository.get_active_skill_file_id(workflow_id, sid)
        if active_skill_file_id:
            self._ensure_skill_leader(workflow_id, active_skill_file_id, preferred_sid=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)
            and self._repository.is_graph_active(workflow_id, current_leader)
        ):
            self._repository.expire_leader(workflow_id)
            return

        if current_leader:
            self._repository.delete_leader(workflow_id)

        new_leader_sid = self._select_graph_leader(workflow_id, preferred_sid=sid)
        if not new_leader_sid:
            self.broadcast_leader_change(workflow_id, None)
            return

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

    def _ensure_skill_leader(self, workflow_id: str, file_id: str, preferred_sid: str | None = None) -> None:
        current_leader = self._repository.get_skill_leader(workflow_id, file_id)
        active_sids = self._repository.get_active_skill_session_sids(workflow_id, file_id)
        if current_leader and self.is_session_active(workflow_id, current_leader):
            if current_leader in active_sids or not active_sids:
                self._repository.expire_skill_leader(workflow_id, file_id)
                return

        if current_leader:
            self._repository.delete_skill_leader(workflow_id, file_id)

        new_leader_sid = self._select_skill_leader(workflow_id, file_id, preferred_sid=preferred_sid)
        if not new_leader_sid:
            self.broadcast_skill_leader_change(workflow_id, file_id, None)
            return

        self._repository.set_skill_leader(workflow_id, file_id, new_leader_sid)
        self.broadcast_skill_leader_change(workflow_id, file_id, new_leader_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") 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 _select_skill_leader(self, workflow_id: str, file_id: str, preferred_sid: str | None = None) -> str | None:
        session_sids = [
            sid
            for sid in self._repository.get_active_skill_session_sids(workflow_id, file_id)
            if self.is_session_active(workflow_id, 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

        try:
            if not self._socketio.manager.is_connected(sid, "/"):
                return False
        except AttributeError:
            return False

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

        if not self._repository.sid_mapping_exists(sid):
            return False

        return True