fix: support multi-worker workflow collaboration (#38242)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Blackoutta <hyytez@gmail.com>
Co-authored-by: Blackoutta <37723456+Blackoutta@users.noreply.github.com>

api/tests/unit_tests/extensions/test_ext_socketio.py

@@ -0,0 +1,33 @@
+import ssl
+
+import socketio
+
+from extensions import ext_socketio
+
+
+def test_socketio_server_uses_redis_manager() -> None:
+    assert isinstance(ext_socketio.sio.manager, socketio.RedisManager)
+
+
+def test_create_socketio_client_manager_uses_pubsub_url_and_prefixed_channel(monkeypatch) -> None:
+    monkeypatch.setattr(ext_socketio.dify_config, "PUBSUB_REDIS_URL", "redis://redis.example.com:6380/3")
+    monkeypatch.setattr(ext_socketio.dify_config, "REDIS_KEY_PREFIX", "tenant-a")
+
+    manager = ext_socketio.create_socketio_client_manager()
+
+    assert manager.redis_url == "redis://redis.example.com:6380/3"
+    assert manager.channel == "tenant-a:socketio"
+
+
+def test_build_redis_options_includes_tls_options_for_rediss(monkeypatch) -> None:
+    monkeypatch.setattr(ext_socketio.dify_config, "REDIS_SSL_CERT_REQS", "CERT_REQUIRED")
+    monkeypatch.setattr(ext_socketio.dify_config, "REDIS_SSL_CA_CERTS", "/ca.pem")
+    monkeypatch.setattr(ext_socketio.dify_config, "REDIS_SSL_CERTFILE", "/cert.pem")
+    monkeypatch.setattr(ext_socketio.dify_config, "REDIS_SSL_KEYFILE", "/key.pem")
+
+    options = ext_socketio._build_redis_options("rediss://redis.example.com:6380/3")
+
+    assert options["ssl_cert_reqs"] == ssl.CERT_REQUIRED
+    assert options["ssl_ca_certs"] == "/ca.pem"
+    assert options["ssl_certfile"] == "/cert.pem"
+    assert options["ssl_keyfile"] == "/key.pem"

api/tests/unit_tests/extensions/test_redis.py

@@ -192,9 +192,13 @@ class TestRedisClientWrapperKeyPrefix:
 
             wrapper.hset("hash:key", "field", "value")
             wrapper.hgetall("hash:key")
+            wrapper.hkeys("hash:key")
+            wrapper.hexists("hash:key", "field")
 
         mock_client.hset.assert_called_once_with("enterprise-a:hash:key", "field", "value")
         mock_client.hgetall.assert_called_once_with("enterprise-a:hash:key")
+        mock_client.hkeys.assert_called_once_with("enterprise-a:hash:key")
+        mock_client.hexists.assert_called_once_with("enterprise-a:hash:key", "field")
 
     def test_wrapper_zadd_prefixes_sorted_set_name(self):
         mock_client = MagicMock()

api/tests/unit_tests/repositories/test_workflow_collaboration_repository.py

@@ -16,14 +16,14 @@ class TestWorkflowCollaborationRepository:
 
     def test_get_sid_mapping_returns_mapping(self, mock_redis: Mock) -> None:
         # Arrange
-        mock_redis.get.return_value = b'{"workflow_id":"wf-1","user_id":"u-1"}'
+        mock_redis.get.return_value = b'{"workflow_id":"wf-1","user_id":"u-1","server_id":"server-1"}'
         repository = WorkflowCollaborationRepository()
 
         # Act
         result = repository.get_sid_mapping("sid-1")
 
         # Assert
-        assert result == {"workflow_id": "wf-1", "user_id": "u-1"}
+        assert result == {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-1"}
 
     def test_list_sessions_filters_invalid_entries(self, mock_redis: Mock) -> None:
         # Arrange
@@ -58,6 +58,7 @@ class TestWorkflowCollaborationRepository:
             "avatar": None,
             "sid": "sid-1",
             "connected_at": 1,
+            "server_id": "server-1",
         }
 
         # Act
@@ -65,11 +66,27 @@ class TestWorkflowCollaborationRepository:
 
         # Assert
         assert mock_redis.hset.called
-        workflow_key, sid, session_json = mock_redis.hset.call_args.args
+        workflow_key, sid, session_json = mock_redis.hset.call_args_list[0].args
         assert workflow_key == "workflow_online_users:wf-1"
         assert sid == "sid-1"
         assert json.loads(session_json)["user_id"] == "u-1"
+        server_sessions_key, server_sid, server_workflow_id = mock_redis.hset.call_args_list[1].args
+        assert server_sessions_key == "ws_server_sessions:server-1"
+        assert server_sid == "sid-1"
+        assert server_workflow_id == "wf-1"
         assert mock_redis.set.called
+        _sid_key, sid_mapping_json = mock_redis.set.call_args.args
+        assert json.loads(sid_mapping_json)["server_id"] == "server-1"
+
+    def test_delete_session_removes_server_session_mapping(self, mock_redis: Mock) -> None:
+        mock_redis.get.return_value = b'{"workflow_id":"wf-1","user_id":"u-1","server_id":"server-1"}'
+        repository = WorkflowCollaborationRepository()
+
+        repository.delete_session("wf-1", "sid-1")
+
+        mock_redis.hdel.assert_any_call("ws_server_sessions:server-1", "sid-1")
+        mock_redis.hdel.assert_any_call("workflow_online_users:wf-1", "sid-1")
+        mock_redis.delete.assert_called_once_with("ws_sid_map:sid-1")
 
     def test_refresh_session_state_expires_keys(self, mock_redis: Mock) -> None:
         # Arrange
@@ -119,3 +136,41 @@ class TestWorkflowCollaborationRepository:
 
         # Assert
         assert result == ["sid-1", "sid-2"]
+
+    def test_refresh_server_heartbeat_sets_ttl(self, mock_redis: Mock) -> None:
+        repository = WorkflowCollaborationRepository()
+
+        repository.refresh_server_heartbeat("server-1")
+
+        mock_redis.set.assert_called_once()
+        key, value = mock_redis.set.call_args.args
+        assert key == "ws_server_heartbeat:server-1"
+        assert value == "1"
+        assert mock_redis.set.call_args.kwargs["ex"] > 0
+
+    def test_server_heartbeat_exists(self, mock_redis: Mock) -> None:
+        mock_redis.exists.return_value = 1
+        repository = WorkflowCollaborationRepository()
+
+        assert repository.server_heartbeat_exists("server-1") is True
+        mock_redis.exists.assert_called_once_with("ws_server_heartbeat:server-1")
+
+    def test_refresh_server_sessions_refreshes_owned_session_ttls(self, mock_redis: Mock) -> None:
+        mock_redis.hgetall.return_value = {b"sid-1": b"wf-1"}
+        mock_redis.exists.return_value = 1
+        repository = WorkflowCollaborationRepository()
+
+        repository.refresh_server_sessions("server-1")
+
+        mock_redis.expire.assert_any_call("workflow_online_users:wf-1", 3600)
+        mock_redis.expire.assert_any_call("ws_sid_map:sid-1", 3600)
+        mock_redis.expire.assert_any_call("ws_server_sessions:server-1", 3600)
+
+    def test_refresh_server_sessions_drops_stale_sid(self, mock_redis: Mock) -> None:
+        mock_redis.hgetall.return_value = {b"sid-stale": b"wf-1"}
+        mock_redis.exists.return_value = 0
+        repository = WorkflowCollaborationRepository()
+
+        repository.refresh_server_sessions("server-1")
+
+        mock_redis.hdel.assert_called_once_with("ws_server_sessions:server-1", "sid-stale")

api/tests/unit_tests/services/test_workflow_collaboration_service.py

@@ -12,7 +12,7 @@ class TestWorkflowCollaborationService:
     def service(self) -> tuple[WorkflowCollaborationService, Mock, Mock]:
         repository = Mock(spec=WorkflowCollaborationRepository)
         socketio = Mock()
-        return WorkflowCollaborationService(repository, socketio), repository, socketio
+        return WorkflowCollaborationService(repository, socketio, server_id="server-1"), repository, socketio
 
     def test_authorize_and_join_workflow_room_returns_leader_status(
         self, service: tuple[WorkflowCollaborationService, Mock, Mock]
@@ -37,6 +37,10 @@ class TestWorkflowCollaborationService:
         # Assert
         assert result == ("u-1", True)
         repository.set_session_info.assert_called_once()
+        session_info = repository.set_session_info.call_args.args[1]
+        assert session_info["server_id"] == "server-1"
+        repository.refresh_server_heartbeat.assert_called_once_with("server-1")
+        socketio.start_background_task.assert_called_once()
         socketio.enter_room.assert_called_once_with("sid-1", "wf-1")
         socketio.emit.assert_called_once_with("status", {"isLeader": True}, room="sid-1")
 
@@ -591,23 +595,47 @@ class TestWorkflowCollaborationService:
 
     def test_is_session_active_guard_branches(self, service: tuple[WorkflowCollaborationService, Mock, Mock]) -> None:
         collaboration_service, repository, socketio = service
-        socketio.manager.is_connected.return_value = True
+        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-1"}
         repository.session_exists.return_value = True
-        repository.sid_mapping_exists.return_value = True
 
         assert collaboration_service.is_session_active("wf-1", "") is False
 
         socketio.manager.is_connected.return_value = False
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
 
+        socketio.manager.is_connected.return_value = True
+        assert collaboration_service.is_session_active("wf-1", "sid-1") is True
+
         socketio.manager.is_connected.side_effect = AttributeError("missing manager")
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
         socketio.manager.is_connected.side_effect = None
 
-        socketio.manager.is_connected.return_value = True
         repository.session_exists.return_value = False
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
 
         repository.session_exists.return_value = True
-        repository.sid_mapping_exists.return_value = False
+        repository.get_sid_mapping.return_value = None
         assert collaboration_service.is_session_active("wf-1", "sid-1") is False
+
+    def test_is_session_active_accepts_remote_session_with_live_server(
+        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
+    ) -> None:
+        collaboration_service, repository, socketio = service
+        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-2"}
+        repository.session_exists.return_value = True
+        repository.server_heartbeat_exists.return_value = True
+        socketio.manager.is_connected.return_value = False
+
+        assert collaboration_service.is_session_active("wf-1", "sid-remote") is True
+        repository.server_heartbeat_exists.assert_called_once_with("server-2")
+
+    def test_is_session_active_rejects_remote_session_with_dead_server(
+        self, service: tuple[WorkflowCollaborationService, Mock, Mock]
+    ) -> None:
+        collaboration_service, repository, socketio = service
+        repository.get_sid_mapping.return_value = {"workflow_id": "wf-1", "user_id": "u-1", "server_id": "server-2"}
+        repository.session_exists.return_value = True
+        repository.server_heartbeat_exists.return_value = False
+        socketio.manager.is_connected.return_value = False
+
+        assert collaboration_service.is_session_active("wf-1", "sid-remote") is False