mirror of
https://github.com/langgenius/dify.git
synced 2026-02-12 04:01:29 -05:00
392 lines
15 KiB
Python
392 lines
15 KiB
Python
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
|