From 0dcd3cd978cff3ecb0bffc700df69f9613a509e5 Mon Sep 17 00:00:00 2001 From: Hamza Rashid Date: Mon, 1 Jun 2026 20:22:11 +0000 Subject: [PATCH] feat(collaboration): reconcile presence via periodic snapshots Reuse the collaboration poll loop to purge expired SQLite connections every 30s and broadcast presence.snapshot per active flow, so clients recover from missed presence.left events. Batch list_users across flow IDs, key connections by connection_id with non-null selected_at, and fold selection into presence.snapshot on the wire and frontend. --- .../langflow/api/utils/collab/connection.py | 77 ++-------- .../langflow/api/v1/collaboration_manager.py | 107 +++++++++++-- .../api/v1/schemas/flow_collaboration.py | 6 +- .../services/collaboration_events/service.py | 4 +- .../services/collaboration_events/sqlite.py | 140 ++++++++++-------- .../unit/api/v1/test_flow_collaboration.py | 21 ++- .../api/v1/test_flow_collaboration_manager.py | 37 ++++- .../services/test_collaboration_events.py | 122 +++++++++++++-- .../flow-collaboration-state.test.ts | 24 ++- .../__tests__/use-flow-collaboration.test.ts | 30 +++- .../hooks/flows/flow-collaboration-state.ts | 24 +-- .../src/hooks/flows/use-flow-collaboration.ts | 9 +- .../src/types/flow-collaboration/index.ts | 7 +- 13 files changed, 398 insertions(+), 210 deletions(-) diff --git a/src/backend/base/langflow/api/utils/collab/connection.py b/src/backend/base/langflow/api/utils/collab/connection.py index 430036bbc2..ab609162cb 100644 --- a/src/backend/base/langflow/api/utils/collab/connection.py +++ b/src/backend/base/langflow/api/utils/collab/connection.py @@ -23,7 +23,6 @@ from langflow.api.utils.collab.operations import ( apply_flow_operation_batch, ) from langflow.api.v1.collaboration_manager import ( - WORKER_ID, CollaborationManager, ensure_collaboration_poll_loop, get_collaboration_manager, @@ -40,7 +39,6 @@ from langflow.api.v1.schemas.flow_collaboration import ( ) from langflow.services.collaboration_events.schemas import ( CollaborationPresenceChange, - CollaborationSelectionTarget, ) from langflow.services.database.models.user.model import UserRead from langflow.services.deps import get_collaboration_events_service, session_scope @@ -146,7 +144,7 @@ class FlowCollaborationConnection: username=self.current_user.username, profile_image=self.current_user.profile_image, ) - snapshot = self._event_service.list_users(self.flow_id) + snapshot = self._event_service.list_users([self.flow_id])[self.flow_id] await self.websocket.send_json( CollaborationSessionReadyMessage( @@ -156,7 +154,6 @@ class FlowCollaborationConnection: ).model_dump(mode="json") ) await self.websocket.send_json(self.manager.presence_snapshot_message(snapshot)) - await self.websocket.send_json(self.manager.selection_snapshot_message(snapshot)) await self._emit_presence_change(presence_change, exclude_connection_id=self.connection_id) @@ -337,28 +334,12 @@ class FlowCollaborationConnection: *, exclude_connection_id: str | None = None, ) -> None: - if change is None: - return - - if change.joined: - message = self.manager.presence_joined_message( - user_id=change.joined.user_id, - username=change.joined.username, - profile_image=change.joined.profile_image, - ) - await self.manager.broadcast_json( - self.flow_id, - message, - exclude_connection_id=exclude_connection_id, - ) - self._publish_presence_joined(change.joined) - - if change.left_user_id: - message = self.manager.presence_left_message(change.left_user_id) - await self.manager.broadcast_json(self.flow_id, message) - self._publish_presence_left(change.left_user_id) - - await self._emit_selection_change(change) + await self.manager.emit_presence_change( + self.flow_id, + change, + self._event_service, + exclude_connection_id=exclude_connection_id, + ) async def _emit_selection_change( self, @@ -366,47 +347,9 @@ class FlowCollaborationConnection: *, exclude_connection_id: str | None = None, ) -> None: - if change is None or change.selection_updated is None: - return - - selected = change.selection_updated.selected - message = self.manager.selection_updated_message(change.selection_updated.user_id, selected) - await self.manager.broadcast_json( + await self.manager.emit_selection_change( self.flow_id, - message, + change, + self._event_service, exclude_connection_id=exclude_connection_id, ) - self._publish_selection_updated(change.selection_updated.user_id, selected) - - def _publish_presence_joined(self, joined: object) -> None: - self._event_service.publish( - self.flow_id, - "presence.joined", - { - "worker_id": WORKER_ID, - "user": { - "user_id": str(joined.user_id), - "username": joined.username, - "profile_image": joined.profile_image, - }, - }, - ) - - def _publish_presence_left(self, user_id: UUID) -> None: - self._event_service.publish( - self.flow_id, - "presence.left", - {"worker_id": WORKER_ID, "user_id": str(user_id)}, - ) - - def _publish_selection_updated( - self, - user_id: UUID, - selected: CollaborationSelectionTarget | None, - ) -> None: - payload: dict[str, object] = { - "worker_id": WORKER_ID, - "user_id": str(user_id), - "selected": {"kind": selected.kind, "id": selected.id} if selected is not None else None, - } - self._event_service.publish(self.flow_id, "selection.updated", payload) diff --git a/src/backend/base/langflow/api/v1/collaboration_manager.py b/src/backend/base/langflow/api/v1/collaboration_manager.py index 6507b59ac1..65a38383f4 100644 --- a/src/backend/base/langflow/api/v1/collaboration_manager.py +++ b/src/backend/base/langflow/api/v1/collaboration_manager.py @@ -23,14 +23,17 @@ from langflow.api.v1.schemas.flow_collaboration import ( CollaborationPresenceLeftMessage, CollaborationPresenceSnapshotMessage, CollaborationPresenceUser, - CollaborationSelectionSnapshotMessage, CollaborationSelectionUpdatedBackplaneEvent, CollaborationSelectionUpdatedMessage, - CollaborationUserSelection, UnsupportedCollaborationBackplaneEventError, parse_collaboration_backplane_event, ) -from langflow.services.collaboration_events.schemas import CollaborationPresenceSnapshot, CollaborationSelectionTarget +from langflow.services.collaboration_events import CollaborationEventService +from langflow.services.collaboration_events.schemas import ( + CollaborationPresenceChange, + CollaborationPresenceSnapshot, + CollaborationSelectionTarget, +) if TYPE_CHECKING: from starlette.websockets import WebSocket @@ -42,6 +45,7 @@ BackplaneEventType = Literal["operation.accepted", "presence.joined", "presence. WORKER_ID = str(uuid.uuid4()) FANNED_REVISION_TTL_SECONDS = 120.0 +PRESENCE_RECONCILE_INTERVAL_SECONDS = 30.0 @dataclass @@ -117,6 +121,7 @@ class CollaborationManager: user_id=user.user_id, username=user.username, profile_image=user.profile_image, + selected=user.selected, ) for user in snapshot.users ] @@ -135,14 +140,6 @@ class CollaborationManager: def presence_left_message(self, user_id: UUID) -> dict[str, Any]: return CollaborationPresenceLeftMessage(user_id=user_id).model_dump(mode="json") - def selection_snapshot_message(self, snapshot: CollaborationPresenceSnapshot) -> dict[str, Any]: - selections = [ - CollaborationUserSelection(user_id=user.user_id, selected=user.selected) - for user in snapshot.users - if user.selected is not None - ] - return CollaborationSelectionSnapshotMessage(selections=selections).model_dump(mode="json") - def selection_updated_message( self, user_id: UUID, @@ -191,6 +188,82 @@ class CollaborationManager: return await handler(flow_id, backplane_event) + async def emit_presence_change( + self, + flow_id: UUID, + change: CollaborationPresenceChange | None, + event_service: CollaborationEventService, + *, + exclude_connection_id: str | None = None, + ) -> None: + if change is None: + return + + if change.joined: + message = self.presence_joined_message( + user_id=change.joined.user_id, + username=change.joined.username, + profile_image=change.joined.profile_image, + ) + await self.broadcast_json( + flow_id, + message, + exclude_connection_id=exclude_connection_id, + ) + event_service.publish( + flow_id, + "presence.joined", + { + "worker_id": WORKER_ID, + "user": { + "user_id": str(change.joined.user_id), + "username": change.joined.username, + "profile_image": change.joined.profile_image, + }, + }, + ) + + if change.left_user_id: + message = self.presence_left_message(change.left_user_id) + await self.broadcast_json(flow_id, message) + event_service.publish( + flow_id, + "presence.left", + {"worker_id": WORKER_ID, "user_id": str(change.left_user_id)}, + ) + + await self.emit_selection_change( + flow_id, + change, + event_service, + exclude_connection_id=exclude_connection_id, + ) + + async def emit_selection_change( + self, + flow_id: UUID, + change: CollaborationPresenceChange | None, + event_service: CollaborationEventService, + *, + exclude_connection_id: str | None = None, + ) -> None: + if change is None or change.selection_updated is None: + return + + selected = change.selection_updated.selected + message = self.selection_updated_message(change.selection_updated.user_id, selected) + await self.broadcast_json( + flow_id, + message, + exclude_connection_id=exclude_connection_id, + ) + payload: dict[str, object] = { + "worker_id": WORKER_ID, + "user_id": str(change.selection_updated.user_id), + "selected": {"kind": selected.kind, "id": selected.id} if selected is not None else None, + } + event_service.publish(flow_id, "selection.updated", payload) + async def _handle_operation_accepted_backplane_event( self, flow_id: UUID, @@ -299,8 +372,9 @@ async def _poll_collaboration_events() -> None: from langflow.services.deps import get_collaboration_events_service manager = get_collaboration_manager() - event_service = get_collaboration_events_service() + event_service: CollaborationEventService = get_collaboration_events_service() cursors: dict[UUID, CollaborationPollCursor] = {} + next_presence_reconcile_at = time.monotonic() while True: flow_ids = manager.active_flow_ids() @@ -308,6 +382,15 @@ async def _poll_collaboration_events() -> None: await asyncio.sleep(0.5) continue + now = time.monotonic() + should_reconcile_presence = now >= next_presence_reconcile_at + if should_reconcile_presence: + next_presence_reconcile_at = now + PRESENCE_RECONCILE_INTERVAL_SECONDS + + snapshots = event_service.list_users(list(flow_ids)) + for flow_id, snapshot in snapshots.items(): + await manager.broadcast_json(flow_id, manager.presence_snapshot_message(snapshot)) + for flow_id in flow_ids: cursor = cursors.get(flow_id) events, new_cursor = event_service.poll(flow_id, cursor=cursor) diff --git a/src/backend/base/langflow/api/v1/schemas/flow_collaboration.py b/src/backend/base/langflow/api/v1/schemas/flow_collaboration.py index 11ae7636f3..fc5d340e6f 100644 --- a/src/backend/base/langflow/api/v1/schemas/flow_collaboration.py +++ b/src/backend/base/langflow/api/v1/schemas/flow_collaboration.py @@ -16,6 +16,7 @@ class CollaborationPresenceUser(BaseModel): user_id: UUID username: str profile_image: str | None = None + selected: CollaborationSelectionTarget | None = None class CollaborationSessionStartMessage(BaseModel): @@ -82,11 +83,6 @@ class CollaborationSelectionUpdateMessage(BaseModel): selected: CollaborationSelectionTarget | None = None -class CollaborationSelectionSnapshotMessage(BaseModel): - type: Literal["selection.snapshot"] = "selection.snapshot" - selections: list[CollaborationUserSelection] - - class CollaborationSelectionUpdatedMessage(BaseModel): type: Literal["selection.updated"] = "selection.updated" user_id: UUID diff --git a/src/backend/base/langflow/services/collaboration_events/service.py b/src/backend/base/langflow/services/collaboration_events/service.py index c6bc29d4bb..9d3def176f 100644 --- a/src/backend/base/langflow/services/collaboration_events/service.py +++ b/src/backend/base/langflow/services/collaboration_events/service.py @@ -79,5 +79,5 @@ class CollaborationEventService(Service, ABC): """Remove one active connection row.""" @abstractmethod - def list_users(self, flow_id: UUID) -> CollaborationPresenceSnapshot: - """Return active deduped users plus effective per-user selections.""" + def list_users(self, flow_ids: list[UUID]) -> dict[UUID, CollaborationPresenceSnapshot]: + """Return active deduped users plus effective per-user selections keyed by flow id.""" diff --git a/src/backend/base/langflow/services/collaboration_events/sqlite.py b/src/backend/base/langflow/services/collaboration_events/sqlite.py index c4d584c755..c26a8f4ad1 100644 --- a/src/backend/base/langflow/services/collaboration_events/sqlite.py +++ b/src/backend/base/langflow/services/collaboration_events/sqlite.py @@ -37,14 +37,13 @@ CREATE INDEX IF NOT EXISTS idx_events_expires ON events(expires_at); CREATE TABLE IF NOT EXISTS connections ( flow_id TEXT NOT NULL, user_id TEXT NOT NULL, - connection_id TEXT NOT NULL, + connection_id TEXT NOT NULL PRIMARY KEY, username TEXT NOT NULL, profile_image TEXT, selected_kind TEXT, selected_id TEXT, - selected_at REAL, - expires_at REAL NOT NULL, - PRIMARY KEY (flow_id, user_id, connection_id) + selected_at REAL NOT NULL, + expires_at REAL NOT NULL ); CREATE INDEX IF NOT EXISTS idx_connections_flow_expires ON connections(flow_id, expires_at); @@ -113,7 +112,7 @@ class SQLiteCollaborationEventService(CollaborationEventService): payload=event_payload, ) expires_at = now + self.TTL_SECONDS - self._purge_expired_locked(now) + self._purge_expired_events_locked(now) self._conn.execute( """ INSERT INTO events @@ -148,7 +147,7 @@ class SQLiteCollaborationEventService(CollaborationEventService): cursor = cursor or CollaborationPollCursor() with self._lock: now = time.time() - self._purge_expired_locked(now) + self._purge_expired_events_locked(now) with self._lock: now = time.time() @@ -219,7 +218,7 @@ class SQLiteCollaborationEventService(CollaborationEventService): self._conn.execute("BEGIN IMMEDIATE") try: now = time.time() - self._purge_expired_locked(now) + self._purge_expired_events_locked(now) before_user = self._effective_user_locked(flow_id_key, user_id_key, now) expires_at = now + self.PRESENCE_TTL_SECONDS self._conn.execute( @@ -227,10 +226,15 @@ class SQLiteCollaborationEventService(CollaborationEventService): INSERT INTO connections ( flow_id, user_id, connection_id, username, profile_image, selected_kind, selected_id, selected_at, expires_at - ) VALUES (?, ?, ?, ?, ?, NULL, NULL, NULL, ?) - ON CONFLICT(flow_id, user_id, connection_id) DO UPDATE SET + ) VALUES (?, ?, ?, ?, ?, NULL, NULL, ?, ?) + ON CONFLICT(connection_id) DO UPDATE SET + flow_id = excluded.flow_id, + user_id = excluded.user_id, username = excluded.username, profile_image = excluded.profile_image, + selected_kind = NULL, + selected_id = NULL, + selected_at = excluded.selected_at, expires_at = excluded.expires_at """, ( @@ -239,6 +243,7 @@ class SQLiteCollaborationEventService(CollaborationEventService): connection_id, username, profile_image, + now, expires_at, ), ) @@ -269,7 +274,7 @@ class SQLiteCollaborationEventService(CollaborationEventService): self._conn.execute("BEGIN IMMEDIATE") try: now = time.time() - self._purge_expired_locked(now) + self._purge_expired_events_locked(now) row = self._conn.execute( """ @@ -351,7 +356,7 @@ class SQLiteCollaborationEventService(CollaborationEventService): self._conn.execute("BEGIN IMMEDIATE") try: now = time.time() - self._purge_expired_locked(now) + self._purge_expired_events_locked(now) row = self._conn.execute( """ @@ -386,21 +391,47 @@ class SQLiteCollaborationEventService(CollaborationEventService): return change - def list_users(self, flow_id: UUID) -> CollaborationPresenceSnapshot: - """Return a snapshot of all active, deduplicated users and their effective selections.""" - flow_id_key = str(flow_id) + def list_users(self, flow_ids: list[UUID]) -> dict[UUID, CollaborationPresenceSnapshot]: + """Return snapshots of active, deduplicated users and their effective selections, grouped by flow id.""" + if not flow_ids: + msg = "flow_ids must not be empty" + raise ValueError(msg) + + flow_id_keys = [str(flow_id) for flow_id in flow_ids] + flow_ids_json = json.dumps(flow_id_keys) with self._lock: - now = time.time() - self._purge_expired_locked(now) - users = self._effective_users_locked(flow_id_key, now) - return CollaborationPresenceSnapshot(users=list(users.values())) + self._conn.execute("BEGIN IMMEDIATE") + try: + now = time.time() + self._purge_expired_events_locked(now) + self._conn.execute( + """ + DELETE FROM connections + WHERE flow_id IN (SELECT value FROM json_each(?)) + AND expires_at < ? + """, + (flow_ids_json, now), + ) + snapshots = self._presence_snapshots_for_flows_locked(flow_id_keys, flow_ids_json, now) + self._conn.execute("COMMIT") + except Exception: + self._conn.execute("ROLLBACK") + raise + + return {flow_id: snapshots.get(flow_id, CollaborationPresenceSnapshot(users=[])) for flow_id in flow_ids} async def teardown(self) -> None: with self._lock: self._conn.close() def _purge_expired_locked(self, now: float) -> None: + self._purge_expired_events_locked(now) + self._purge_expired_connections_locked(now) + + def _purge_expired_events_locked(self, now: float) -> None: self._conn.execute("DELETE FROM events WHERE expires_at < ?", (now,)) + + def _purge_expired_connections_locked(self, now: float) -> None: self._conn.execute("DELETE FROM connections WHERE expires_at < ?", (now,)) def _enforce_per_flow_cap_locked(self, flow_id: str) -> None: @@ -417,63 +448,50 @@ class SQLiteCollaborationEventService(CollaborationEventService): (flow_id, self.MAX_EVENTS_PER_FLOW), ) - def _effective_users_locked( + def _presence_snapshots_for_flows_locked( self, - flow_id_key: str, + flow_id_keys: list[str], + flow_ids_json: str, now: float, - ) -> dict[UUID, CollaborationPresenceConnectionUser]: - """Compute the effective deduplicated user presence state. - - Groups all non-expired connections by user_id. The effective selection - for each user is taken from the connection that most recently updated - its selection. - """ + ) -> dict[UUID, CollaborationPresenceSnapshot]: rows = self._conn.execute( """ - SELECT user_id, connection_id, username, profile_image, - selected_kind, selected_id, selected_at - FROM connections - WHERE flow_id = ? + SELECT + flow_id, + user_id, + username, + profile_image, + selected_kind, + selected_id + FROM connections AS c + WHERE flow_id IN (SELECT value FROM json_each(?)) AND expires_at >= ? - ORDER BY user_id ASC, selected_at DESC + ORDER BY flow_id ASC, user_id ASC, selected_at DESC, connection_id ASC """, - (flow_id_key, now), + (flow_ids_json, now), ).fetchall() - users: dict[UUID, CollaborationPresenceConnectionUser] = {} - best_selection: dict[UUID, tuple[float, CollaborationSelectionTarget | None]] = {} - + snapshots = {UUID(flow_id_key): CollaborationPresenceSnapshot(users=[]) for flow_id_key in flow_id_keys} + seen_users: set[tuple[UUID, UUID]] = set() for row in rows: - user_id = UUID(row[0]) - username = row[2] - profile_image = row[3] - selected_kind = row[4] - selected_id = row[5] - selected_at = row[6] + flow_id = UUID(row[0]) + user_id = UUID(row[1]) + user_key = (flow_id, user_id) + if user_key in seen_users: + continue + seen_users.add(user_key) - if user_id not in users: - users[user_id] = CollaborationPresenceConnectionUser( + selected = self._selection_target(row[4], row[5]) + snapshots[flow_id].users.append( + CollaborationPresenceConnectionUser( user_id=user_id, - username=username, - profile_image=profile_image, + username=row[2], + profile_image=row[3], + selected=selected, ) - - if selected_at is not None: - selected = self._selection_target(selected_kind, selected_id) - current = best_selection.get(user_id) - if current is None or selected_at > current[0]: - best_selection[user_id] = (selected_at, selected) - - for user_id, (_, selected) in best_selection.items(): - user = users[user_id] - users[user_id] = CollaborationPresenceConnectionUser( - user_id=user.user_id, - username=user.username, - profile_image=user.profile_image, - selected=selected, ) - return users + return snapshots def _effective_user_locked( self, diff --git a/src/backend/tests/unit/api/v1/test_flow_collaboration.py b/src/backend/tests/unit/api/v1/test_flow_collaboration.py index 93950ea8c9..8f4157be49 100644 --- a/src/backend/tests/unit/api/v1/test_flow_collaboration.py +++ b/src/backend/tests/unit/api/v1/test_flow_collaboration.py @@ -99,11 +99,10 @@ def _receive_message_type(ws, expected_type: str) -> dict: return message -def _receive_session_bootstrap(ws) -> tuple[dict, dict, dict]: +def _receive_session_bootstrap(ws) -> tuple[dict, dict]: ready = _receive_message_type(ws, "session.ready") presence = _receive_message_type(ws, "presence.snapshot") - selection = _receive_message_type(ws, "selection.snapshot") - return ready, presence, selection + return ready, presence def _close_websocket_cleanly(ws) -> None: @@ -262,7 +261,7 @@ async def test_collab_session_ready(client: AsyncClient, logged_in_headers): def _assert(ws) -> None: ws.send_json({"type": "session.start"}) - ready, presence, selection = _receive_session_bootstrap(ws) + ready, presence = _receive_session_bootstrap(ws) assert ready["type"] == "session.ready" assert ready["current_revision"] == 0 assert ready["flow_id"] == str(flow_id) @@ -270,8 +269,6 @@ async def test_collab_session_ready(client: AsyncClient, logged_in_headers): assert presence["type"] == "presence.snapshot" assert len(presence["users"]) == 1 assert presence["users"][0]["username"] == "activeuser" - assert selection["type"] == "selection.snapshot" - assert selection["selections"] == [] await _run_websocket_test(app, flow_id, token, _assert) @@ -283,7 +280,7 @@ async def test_operation_submit_accepted_increments_revision(client: AsyncClient def _submit(ws) -> None: ws.send_json({"type": "session.start"}) - ready, _, _ = _receive_session_bootstrap(ws) + ready, _ = _receive_session_bootstrap(ws) updated = copy.deepcopy(NODE_A) updated["position"] = {"x": 50, "y": 50} ws.send_json( @@ -362,7 +359,7 @@ async def test_stale_revision_rejected(client: AsyncClient, logged_in_headers): def _submit(ws) -> None: ws.send_json({"type": "session.start"}) - ready, _, _ = _receive_session_bootstrap(ws) + ready, _ = _receive_session_bootstrap(ws) ws.send_json( { "type": "operation.submit", @@ -425,7 +422,7 @@ async def test_invalid_edge_rejected_without_revision_change(client: AsyncClient def _submit(ws) -> None: ws.send_json({"type": "session.start"}) - ready, _, _ = _receive_session_bootstrap(ws) + ready, _ = _receive_session_bootstrap(ws) ws.send_json( { "type": "operation.submit", @@ -458,7 +455,7 @@ async def test_operation_broadcast_to_peer_socket(client: AsyncClient, logged_in def _peers(ws_a, ws_b) -> None: ws_a.send_json({"type": "session.start"}) - ready_a, _, _ = _receive_session_bootstrap(ws_a) + ready_a, _ = _receive_session_bootstrap(ws_a) ws_b.send_json({"type": "session.start"}) _receive_session_bootstrap(ws_b) @@ -490,12 +487,12 @@ async def test_presence_shows_each_user_once_with_two_tabs(client: AsyncClient, def _peers(ws_a, ws_b) -> None: ws_a.send_json({"type": "session.start"}) - _, presence_a, _ = _receive_session_bootstrap(ws_a) + _, presence_a = _receive_session_bootstrap(ws_a) assert len(presence_a["users"]) == 1 assert presence_a["users"][0]["username"] == "activeuser" ws_b.send_json({"type": "session.start"}) - _, presence_b, _ = _receive_session_bootstrap(ws_b) + _, presence_b = _receive_session_bootstrap(ws_b) assert len(presence_b["users"]) == 1 await _run_dual_websocket_test(app, flow_id, token, _peers) diff --git a/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py b/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py index ebfaf68450..a1e678ec07 100644 --- a/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py +++ b/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock from uuid import uuid4 import pytest @@ -10,6 +10,7 @@ from langflow.api.v1 import collaboration_manager as collaboration_manager_modul from langflow.api.v1.collaboration_manager import CollaborationManager from langflow.services.collaboration_events.schemas import ( CollaborationEvent, + CollaborationPresenceChange, CollaborationPresenceConnectionUser, CollaborationPresenceSnapshot, ) @@ -166,7 +167,7 @@ async def test_handle_backplane_event_ignores_unknown_type(manager, flow_id, use @pytest.mark.asyncio -async def test_presence_snapshot_and_selection_snapshot_messages(manager, user_a): +async def test_presence_snapshot_includes_selection(manager, user_a): snapshot = _snapshot( CollaborationPresenceConnectionUser( user_id=user_a, @@ -177,12 +178,36 @@ async def test_presence_snapshot_and_selection_snapshot_messages(manager, user_a ) presence = manager.presence_snapshot_message(snapshot) - selection = manager.selection_snapshot_message(snapshot) assert presence["type"] == "presence.snapshot" - assert len(presence["users"]) == 1 - assert selection["type"] == "selection.snapshot" - assert selection["selections"] == [{"user_id": str(user_a), "selected": {"kind": "node", "id": "node-1"}}] + assert presence["users"] == [ + { + "user_id": str(user_a), + "username": "alice", + "profile_image": None, + "selected": {"kind": "node", "id": "node-1"}, + } + ] + + +@pytest.mark.asyncio +async def test_emit_presence_change_broadcasts_and_publishes_left(manager, flow_id, user_a, user_b): + await _register(manager, flow_id, user_a, "alice") + peer_id = await _register(manager, flow_id, user_b, "bob") + peer_ws = manager._rooms[flow_id][peer_id].websocket + event_service = Mock() + + await manager.emit_presence_change(flow_id, CollaborationPresenceChange(left_user_id=user_a), event_service) + + peer_ws.send_json.assert_called_once() + payload = peer_ws.send_json.call_args.args[0] + assert payload["type"] == "presence.left" + assert payload["user_id"] == str(user_a) + event_service.publish.assert_called_once_with( + flow_id, + "presence.left", + {"worker_id": collaboration_manager_module.WORKER_ID, "user_id": str(user_a)}, + ) @pytest.mark.asyncio diff --git a/src/backend/tests/unit/services/test_collaboration_events.py b/src/backend/tests/unit/services/test_collaboration_events.py index 9717bb01a3..5e3539e109 100644 --- a/src/backend/tests/unit/services/test_collaboration_events.py +++ b/src/backend/tests/unit/services/test_collaboration_events.py @@ -180,6 +180,11 @@ def test_presence_change_fields_are_mutually_exclusive(): CollaborationPresenceChange(left_user_id=user_id, selection_updated=selection) +def test_list_users_requires_flow_ids(svc: SQLiteCollaborationEventService): + with pytest.raises(ValueError, match="flow_ids"): + svc.list_users([]) + + def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventService, flow_id: UUID): user_id = uuid4() conn_a = "conn-a" @@ -196,7 +201,7 @@ def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventServic assert joined.joined is not None assert joined.joined.username == "alice" - snapshot = svc.list_users(flow_id) + snapshot = svc.list_users([flow_id])[flow_id] assert len(snapshot.users) == 1 second = svc.add_connection( @@ -207,7 +212,7 @@ def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventServic profile_image=None, ) assert second is None - assert len(svc.list_users(flow_id).users) == 1 + assert len(svc.list_users([flow_id])[flow_id].users) == 1 change = svc.update_connection( flow_id=flow_id, @@ -218,20 +223,20 @@ def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventServic assert change.selection_updated is not None assert change.selection_updated.selected == CollaborationSelectionTarget(kind="node", id="node-1") - snapshot = svc.list_users(flow_id) + snapshot = svc.list_users([flow_id])[flow_id] assert snapshot.users[0].selected == CollaborationSelectionTarget(kind="node", id="node-1") left = svc.remove_connection(flow_id=flow_id, connection_id=conn_a) assert left is not None assert left.left_user_id is None assert left.selection_updated == CollaborationUserSelection(user_id=user_id, selected=None) - assert len(svc.list_users(flow_id).users) == 1 + assert len(svc.list_users([flow_id])[flow_id].users) == 1 final = svc.remove_connection(flow_id=flow_id, connection_id=conn_b) assert final is not None assert final.left_user_id == user_id assert final.selection_updated is None - assert svc.list_users(flow_id).users == [] + assert svc.list_users([flow_id])[flow_id].users == [] def test_presence_ttl_cleanup(svc: SQLiteCollaborationEventService, flow_id: UUID): @@ -244,12 +249,84 @@ def test_presence_ttl_cleanup(svc: SQLiteCollaborationEventService, flow_id: UUI username="alice", profile_image=None, ) - assert len(svc.list_users(flow_id).users) == 1 + assert len(svc.list_users([flow_id])[flow_id].users) == 1 time.sleep(0.15) svc.cleanup() - assert svc.list_users(flow_id).users == [] + assert svc.list_users([flow_id])[flow_id].users == [] + + +def test_list_users_removes_expired_connections(svc: SQLiteCollaborationEventService, flow_id: UUID): + svc.PRESENCE_TTL_SECONDS = 0.1 + user_id = uuid4() + svc.add_connection( + flow_id=flow_id, + user_id=user_id, + connection_id="conn-1", + username="alice", + profile_image=None, + ) + + time.sleep(0.15) + + snapshots = svc.list_users([flow_id]) + assert snapshots[flow_id].users == [] + assert svc.list_users([flow_id])[flow_id].users == [] + + +def test_poll_does_not_consume_expired_presence_before_list_users( + svc: SQLiteCollaborationEventService, + flow_id: UUID, +): + svc.PRESENCE_TTL_SECONDS = 0.1 + user_id = uuid4() + svc.add_connection( + flow_id=flow_id, + user_id=user_id, + connection_id="conn-1", + username="alice", + profile_image=None, + ) + + time.sleep(0.15) + svc.publish(flow_id, "operation.accepted", {"revision": 1}) + svc.poll(flow_id) + + snapshots = svc.list_users([flow_id]) + assert snapshots[flow_id].users == [] + + +def test_list_users_batches_flow_snapshots( + svc: SQLiteCollaborationEventService, + flow_id: UUID, + other_flow_id: UUID, +): + svc.PRESENCE_TTL_SECONDS = 0.1 + expired_user_id = uuid4() + active_user_id = uuid4() + svc.add_connection( + flow_id=flow_id, + user_id=expired_user_id, + connection_id="expired-conn", + username="expired", + profile_image=None, + ) + svc.PRESENCE_TTL_SECONDS = 30.0 + svc.add_connection( + flow_id=other_flow_id, + user_id=active_user_id, + connection_id="active-conn", + username="active", + profile_image=None, + ) + + time.sleep(0.15) + + snapshots = svc.list_users([flow_id, other_flow_id]) + assert snapshots[flow_id].users == [] + assert len(snapshots[other_flow_id].users) == 1 + assert snapshots[other_flow_id].users[0].user_id == active_user_id def test_effective_selection_uses_latest_connection(svc: SQLiteCollaborationEventService, flow_id: UUID): @@ -280,10 +357,37 @@ def test_effective_selection_uses_latest_connection(svc: SQLiteCollaborationEven selected=CollaborationSelectionTarget(kind="edge", id="new"), ) - snapshot = svc.list_users(flow_id) + snapshot = svc.list_users([flow_id])[flow_id] assert snapshot.users[0].selected == CollaborationSelectionTarget(kind="edge", id="new") +def test_initial_connection_counts_as_currently_unselected(svc: SQLiteCollaborationEventService, flow_id: UUID): + user_id = uuid4() + svc.add_connection( + flow_id=flow_id, + user_id=user_id, + connection_id="conn-a", + username="alice", + profile_image=None, + ) + svc.update_connection( + flow_id=flow_id, + connection_id="conn-a", + selected=CollaborationSelectionTarget(kind="node", id="old"), + ) + time.sleep(0.01) + svc.add_connection( + flow_id=flow_id, + user_id=user_id, + connection_id="conn-b", + username="alice", + profile_image=None, + ) + + snapshot = svc.list_users([flow_id])[flow_id] + assert snapshot.users[0].selected is None + + def test_cross_worker_presence_visibility(tmp_path, flow_id: UUID): shared = tmp_path / "shared" user_id = uuid4() @@ -299,6 +403,6 @@ def test_cross_worker_presence_visibility(tmp_path, flow_id: UUID): profile_image=None, ) - snapshot = worker_b.list_users(flow_id) + snapshot = worker_b.list_users([flow_id])[flow_id] assert len(snapshot.users) == 1 assert snapshot.users[0].username == "bob" diff --git a/src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts b/src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts index 51c8cc59a0..746a705b61 100644 --- a/src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts +++ b/src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts @@ -2,8 +2,8 @@ import { applyPresenceJoined, applyPresenceLeft, applyPresenceSnapshot, - applySelectionSnapshot, applySelectionUpdated, + selectionsFromPresenceSnapshot, } from "@/hooks/flows/flow-collaboration-state"; describe("flow-collaboration-state", () => { @@ -11,7 +11,13 @@ describe("flow-collaboration-state", () => { expect( applyPresenceSnapshot( [{ user_id: "old", username: "old-user" }], - [{ user_id: "new", username: "new-user" }], + [ + { + user_id: "new", + username: "new-user", + selected: { kind: "node", id: "n1" }, + }, + ], ), ).toEqual([{ user_id: "new", username: "new-user" }]); }); @@ -51,12 +57,16 @@ describe("flow-collaboration-state", () => { ).toEqual([{ user_id: "user-2", username: "bob" }]); }); - it("should replace selections on selection.snapshot", () => { + it("should derive selections from presence.snapshot users", () => { expect( - applySelectionSnapshot( - [{ user_id: "user-1", selected: { kind: "node", id: "n1" } }], - [{ user_id: "user-2", selected: { kind: "edge", id: "e1" } }], - ), + selectionsFromPresenceSnapshot([ + { user_id: "user-1", username: "ana", selected: null }, + { + user_id: "user-2", + username: "bob", + selected: { kind: "edge", id: "e1" }, + }, + ]), ).toEqual([{ user_id: "user-2", selected: { kind: "edge", id: "e1" } }]); }); diff --git a/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration.test.ts b/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration.test.ts index a1ba0555fb..952f0a0402 100644 --- a/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration.test.ts +++ b/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration.test.ts @@ -140,13 +140,10 @@ async function connectSession(currentRevision = 0) { user_id: "user-1", username: "ana", profile_image: "Space/046-rocket.svg", + selected: null, }, ], }); - socket.triggerMessage({ - type: "selection.snapshot", - selections: [], - }); }); return socket; @@ -371,6 +368,26 @@ describe("useFlowCollaboration", () => { const { result } = await mountHook({ flowId: "flow-1" }); await connectSession(0); + await act(async () => { + latestSocket().triggerMessage({ + type: "presence.snapshot", + users: [ + { + user_id: "user-1", + username: "ana", + selected: { kind: "edge", id: "edge-1" }, + }, + ], + }); + }); + + expect(result.current.users).toEqual([ + { user_id: "user-1", username: "ana" }, + ]); + expect(result.current.selections).toEqual([ + { user_id: "user-1", selected: { kind: "edge", id: "edge-1" } }, + ]); + await act(async () => { latestSocket().triggerMessage({ type: "selection.updated", @@ -380,6 +397,7 @@ describe("useFlowCollaboration", () => { }); expect(result.current.selections).toEqual([ + { user_id: "user-1", selected: { kind: "edge", id: "edge-1" } }, { user_id: "user-2", selected: { kind: "node", id: "node-1" } }, ]); @@ -391,7 +409,9 @@ describe("useFlowCollaboration", () => { }); }); - expect(result.current.selections).toEqual([]); + expect(result.current.selections).toEqual([ + { user_id: "user-1", selected: { kind: "edge", id: "edge-1" } }, + ]); }); it("should send selection.update when sendSelectionUpdate is called", async () => { diff --git a/src/frontend/src/hooks/flows/flow-collaboration-state.ts b/src/frontend/src/hooks/flows/flow-collaboration-state.ts index d90b52bb91..ac400d0c02 100644 --- a/src/frontend/src/hooks/flows/flow-collaboration-state.ts +++ b/src/frontend/src/hooks/flows/flow-collaboration-state.ts @@ -8,22 +8,23 @@ export function applyPresenceSnapshot( _currentUsers: CollaborationPresenceUser[], users: CollaborationPresenceUser[], ): CollaborationPresenceUser[] { - return users.map((user) => ({ ...user })); + return users.map(({ selected: _selected, ...user }) => ({ ...user })); } export function applyPresenceJoined( currentUsers: CollaborationPresenceUser[], user: CollaborationPresenceUser, ): CollaborationPresenceUser[] { + const { selected: _selected, ...presenceUser } = user; const existingIndex = currentUsers.findIndex( - (entry) => entry.user_id === user.user_id, + (entry) => entry.user_id === presenceUser.user_id, ); if (existingIndex === -1) { - return [...currentUsers, { ...user }]; + return [...currentUsers, { ...presenceUser }]; } const nextUsers = [...currentUsers]; - nextUsers[existingIndex] = { ...user }; + nextUsers[existingIndex] = { ...presenceUser }; return nextUsers; } @@ -34,14 +35,15 @@ export function applyPresenceLeft( return currentUsers.filter((user) => user.user_id !== userId); } -export function applySelectionSnapshot( - _currentSelections: CollaborationUserSelection[], - selections: CollaborationUserSelection[], +export function selectionsFromPresenceSnapshot( + users: CollaborationPresenceUser[], ): CollaborationUserSelection[] { - return selections.map((selection) => ({ - user_id: selection.user_id, - selected: selection.selected ? { ...selection.selected } : null, - })); + return users + .filter((user) => user.selected != null) + .map((user) => ({ + user_id: user.user_id, + selected: user.selected ? { ...user.selected } : null, + })); } export function applySelectionUpdated( diff --git a/src/frontend/src/hooks/flows/use-flow-collaboration.ts b/src/frontend/src/hooks/flows/use-flow-collaboration.ts index 50f8822c55..23f6ed5ada 100644 --- a/src/frontend/src/hooks/flows/use-flow-collaboration.ts +++ b/src/frontend/src/hooks/flows/use-flow-collaboration.ts @@ -3,8 +3,8 @@ import { applyPresenceJoined, applyPresenceLeft, applyPresenceSnapshot, - applySelectionSnapshot, applySelectionUpdated, + selectionsFromPresenceSnapshot, } from "@/hooks/flows/flow-collaboration-state"; import { buildFlowCollaborationWebSocketUrl } from "@/hooks/flows/flow-collaboration-url"; import type { @@ -203,6 +203,7 @@ export function useFlowCollaboration({ setUsers((currentUsers) => applyPresenceSnapshot(currentUsers, message.users), ); + setSelections(selectionsFromPresenceSnapshot(message.users)); return; } case "presence.joined": { @@ -222,12 +223,6 @@ export function useFlowCollaboration({ ); return; } - case "selection.snapshot": { - setSelections((currentSelections) => - applySelectionSnapshot(currentSelections, message.selections), - ); - return; - } case "selection.updated": { setSelections((currentSelections) => applySelectionUpdated( diff --git a/src/frontend/src/types/flow-collaboration/index.ts b/src/frontend/src/types/flow-collaboration/index.ts index 3ecfc04aca..0f74e13fe2 100644 --- a/src/frontend/src/types/flow-collaboration/index.ts +++ b/src/frontend/src/types/flow-collaboration/index.ts @@ -8,6 +8,7 @@ export type CollaborationPresenceUser = { user_id: string; username: string; profile_image?: string | null; + selected?: CollaborationSelectionTarget | null; }; export type CollaborationSessionStartMessage = { @@ -86,11 +87,6 @@ export type CollaborationSelectionUpdateMessage = { selected: CollaborationSelectionTarget | null; }; -export type CollaborationSelectionSnapshotMessage = { - type: "selection.snapshot"; - selections: CollaborationUserSelection[]; -}; - export type CollaborationSelectionUpdatedMessage = { type: "selection.updated"; user_id: string; @@ -111,7 +107,6 @@ export type CollaborationServerMessage = | CollaborationPresenceSnapshotMessage | CollaborationPresenceJoinedMessage | CollaborationPresenceLeftMessage - | CollaborationSelectionSnapshotMessage | CollaborationSelectionUpdatedMessage | CollaborationMessageErrorMessage;