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;