establish websocket connection

api/app.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import logging
 
 
 def is_db_command():
@@ -33,9 +34,10 @@ else:
         psycogreen.gevent.patch_psycopg()
 
     from app_factory import create_app
+    from extensions.ext_socketio import ext_socketio
 
     app = create_app()
     celery = app.extensions["celery"]
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5001)
+    ext_socketio.run(app, host="0.0.0.0", port=5001, debug=True)

api/app_factory.py

@@ -57,6 +57,7 @@ def initialize_extensions(app: DifyApp):
         ext_request_logging,
         ext_sentry,
         ext_set_secretkey,
+        ext_socketio,
         ext_storage,
         ext_timezone,
         ext_warnings,
@@ -85,6 +86,7 @@ def initialize_extensions(app: DifyApp):
         ext_commands,
         ext_otel,
         ext_request_logging,
+        ext_socketio,
     ]
     for ext in extensions:
         short_name = ext.__name__.split(".")[-1]

api/controllers/console/app/online_user.py

@@ -0,0 +1,60 @@
+import json
+
+from flask_login import current_user, login_required
+from flask import request
+
+from extensions.ext_socketio import ext_socketio
+from extensions.ext_redis import redis_client
+
+@ext_socketio.on('user_connect')
+@login_required
+def handle_user_connect(data):
+    """
+    Handle user connect event, check login and get user info.
+    """
+    
+    sid = request.sid
+    workflow_id = data.get('workflow_id')
+
+    old_info_json = redis_client.hget(f"workflow_online_users:{workflow_id}", current_user.id)
+    if old_info_json:
+        old_info = json.loads(old_info_json)
+        old_sid = old_info.get("sid")
+        if old_sid and old_sid != sid:
+            ext_socketio.server.disconnect(sid=old_sid)
+
+    user_info = {
+        "user_id": current_user.id,
+        "username": getattr(current_user, "username", ""),
+        "avatar": getattr(current_user, "avatar", ""),
+        "sid": sid
+    } 
+
+    redis_client.hset(
+        f"workflow_online_users:{workflow_id}",
+        current_user.id,
+        json.dumps(user_info)
+    )
+
+    redis_client.set(
+        f"ws_sid_map:{sid}",
+        json.dumps({
+            "workflow_id": workflow_id,
+            "user_id": current_user.id
+        })
+    )
+    return {'msg': 'connected', 'user_id': current_user.id, 'sid': sid}
+
+@ext_socketio.on('disconnect')
+def handle_disconnect():
+    """
+    Handle user disconnect event, remove user from workflow's online user list.
+    """
+    sid = request.sid
+    mapping = redis_client.get(f"ws_sid_map:{sid}")
+    if mapping:
+        data = json.loads(mapping)
+        workflow_id = data["workflow_id"]
+        user_id = data["user_id"]   
+        redis_client.hdel(f"workflow_online_users:{workflow_id}", user_id)
+        redis_client.delete(f"ws_sid_map:{sid}")

api/docker/entrypoint.sh

@@ -28,12 +28,13 @@ elif [[ "${MODE}" == "beat" ]]; then
   exec celery -A app.celery beat --loglevel ${LOG_LEVEL:-INFO}
 else
   if [[ "${DEBUG}" == "true" ]]; then
+  # TODO: add socketio support
     exec flask run --host=${DIFY_BIND_ADDRESS:-0.0.0.0} --port=${DIFY_PORT:-5001} --debug
   else
     exec gunicorn \
       --bind "${DIFY_BIND_ADDRESS:-0.0.0.0}:${DIFY_PORT:-5001}" \
       --workers ${SERVER_WORKER_AMOUNT:-1} \
-      --worker-class ${SERVER_WORKER_CLASS:-gevent} \
+      --worker-class ${SERVER_WORKER_CLASS:-geventwebsocket.gunicorn.workers.GeventWebSocketWorker} \
       --worker-connections ${SERVER_WORKER_CONNECTIONS:-10} \
       --timeout ${GUNICORN_TIMEOUT:-200} \
       app:app

api/extensions/ext_socketio.py

@@ -0,0 +1,10 @@
+from flask_socketio import SocketIO
+
+from configs import dify_config
+from dify_app import DifyApp
+
+ext_socketio = SocketIO()
+
+
+def init_app(app: DifyApp):
+    ext_socketio.init_app(app, async_mode='gevent', cors_allowed_origins=dify_config.CONSOLE_CORS_ALLOW_ORIGINS) 

api/pyproject.toml

@@ -19,8 +19,10 @@ dependencies = [
     "flask-login~=0.6.3",
     "flask-migrate~=4.0.7",
     "flask-restful~=0.3.10",
+    "flask-socketio~=5.5.1",
     "flask-sqlalchemy~=3.1.1",
     "gevent~=24.11.1",
+    "gevent-websocket~=0.10.1",
     "gmpy2~=2.2.1",
     "google-api-core==2.18.0",
     "google-api-python-client==2.90.0",