diff --git a/src/backend/base/langflow/api/utils/collab/connection.py b/src/backend/base/langflow/api/utils/collab/connection.py index c6d5f2743f..f857296055 100644 --- a/src/backend/base/langflow/api/utils/collab/connection.py +++ b/src/backend/base/langflow/api/utils/collab/connection.py @@ -32,6 +32,7 @@ from langflow.api.v1.schemas.flow_collaboration import ( CollaborationOperationBroadcastMessage, CollaborationOperationRejectedMessage, CollaborationOperationSubmitMessage, + CollaborationSelectionUpdateMessage, CollaborationSessionReadyMessage, CollaborationSessionStartMessage, CollaborationUnknownMessageError, @@ -99,6 +100,10 @@ class FlowCollaborationConnection: await self._handle_operation_submit(raw) continue + if msg_type == "selection.update": + await self._handle_selection_update(raw) + continue + await self.websocket.send_json(CollaborationUnknownMessageError().model_dump(mode="json")) async def _handle_session_start(self, raw: Any, msg_type: str | None) -> None: @@ -120,6 +125,8 @@ class FlowCollaborationConnection: ) raise _CollaborationConnectionClosedError from exc + user_was_already_in_room = self.current_user.id in self.manager.all_users(self.flow_id, as_dict=True) + self.connection_id = await self.manager.register( websocket=self.websocket, flow_id=self.flow_id, @@ -135,10 +142,23 @@ class FlowCollaborationConnection: connection_id=self.connection_id, flow_id=self.flow_id, current_revision=self.starting_revision, - users=self.manager.all_users(self.flow_id), ).model_dump(mode="json") ) - await self._publish_presence(exclude_connection_id=self.connection_id) + await self.websocket.send_json(self.manager.presence_snapshot_message(self.flow_id)) + await self.websocket.send_json(self.manager.selection_snapshot_message(self.flow_id)) + + if not user_was_already_in_room: + await self.manager.broadcast_json( + self.flow_id, + self.manager.presence_joined_message( + user_id=self.current_user.id, + username=self.current_user.username, + profile_image=self.current_user.profile_image, + ), + exclude_connection_id=self.connection_id, + ) + + self._publish_presence() async def _handle_operation_submit(self, raw: Any) -> None: try: @@ -262,15 +282,32 @@ class FlowCollaborationConnection: } get_collaboration_events_service().publish(self.flow_id, "operation.accepted", event_payload) - async def _publish_presence(self, *, exclude_connection_id: str | None = None) -> None: - payload = self.manager.presence_payload(self.flow_id) - get_collaboration_events_service().publish(self.flow_id, "presence.updated", payload) + async def _handle_selection_update(self, raw: Any) -> None: + try: + update = CollaborationSelectionUpdateMessage.model_validate(raw) + except ValidationError as exc: + await self.websocket.send_json( + CollaborationUnknownMessageError( + detail=f"Invalid selection.update payload: {exc}", + ).model_dump(mode="json") + ) + return + + updated = self.manager.set_user_selection( + self.flow_id, + self.current_user.id, + update.selected, + ) await self.manager.broadcast_json( self.flow_id, - self.manager.presence_message(self.flow_id), - exclude_connection_id=exclude_connection_id, + updated, + exclude_connection_id=self.connection_id, ) + def _publish_presence(self) -> None: + payload = self.manager.presence_payload(self.flow_id) + get_collaboration_events_service().publish(self.flow_id, "presence.roster", payload) + async def _presence_heartbeat(self) -> None: while True: await asyncio.sleep(_PRESENCE_HEARTBEAT_SECONDS) @@ -282,7 +319,7 @@ class FlowCollaborationConnection: await self._ensure_active_read_access() except _CollaborationConnectionClosedError: return - await self._publish_presence() + self._publish_presence() async def _cleanup(self) -> None: if self._presence_task is not None: @@ -293,8 +330,20 @@ class FlowCollaborationConnection: pass if self.connection_id is not None: - await self.manager.unregister(self.flow_id, self.connection_id) + conn = await self.manager.unregister(self.flow_id, self.connection_id) + + if conn is not None: + selection_update = self.manager.clear_user_selection(self.flow_id, conn.user_id) + if selection_update is not None: + await self.manager.broadcast_json(self.flow_id, selection_update) + + if conn.user_id not in self.manager.all_users(self.flow_id, as_dict=True): + await self.manager.broadcast_json( + self.flow_id, + self.manager.presence_left_message(conn.user_id), + ) + try: - await self._publish_presence() + self._publish_presence() except Exception as exc: # noqa: BLE001 - cleanup presence publish is best-effort. await logger.adebug("Failed to publish collaboration presence during cleanup: %s", exc) diff --git a/src/backend/base/langflow/api/v1/collaboration_manager.py b/src/backend/base/langflow/api/v1/collaboration_manager.py index c29eb0f2e2..75544b37e9 100644 --- a/src/backend/base/langflow/api/v1/collaboration_manager.py +++ b/src/backend/base/langflow/api/v1/collaboration_manager.py @@ -17,8 +17,15 @@ from langflow.api.v1.schemas.flow_collaboration import ( CollaborationOperationAcceptedBackplaneEvent, CollaborationOperationBroadcastMessage, CollaborationPresenceEventPayload, - CollaborationPresenceUpdatedBackplaneEvent, + CollaborationPresenceJoinedMessage, + CollaborationPresenceLeftMessage, + CollaborationPresenceRosterBackplaneEvent, + CollaborationPresenceSnapshotMessage, CollaborationPresenceUser, + CollaborationSelectionSnapshotMessage, + CollaborationSelectionTarget, + CollaborationSelectionUpdatedMessage, + CollaborationUserSelection, UnsupportedCollaborationBackplaneEventError, parse_collaboration_backplane_event, ) @@ -55,6 +62,7 @@ class CollaborationManager: def __init__(self) -> None: self._rooms: defaultdict[UUID, dict[str, FlowConnection]] = defaultdict(dict) self._remote_rosters: defaultdict[UUID, dict[str, _RemotePresenceRoster]] = defaultdict(dict) + self._selections: defaultdict[UUID, dict[UUID, CollaborationSelectionTarget | None]] = defaultdict(dict) self._fanned_revisions: dict[tuple[UUID, int], float] = {} # Protect local room membership while registering/unregistering and snapshotting broadcasts. self._lock = asyncio.Lock() @@ -92,29 +100,67 @@ class CollaborationManager: self._rooms[flow_id][connection_id] = conn return connection_id - async def unregister(self, flow_id: UUID, connection_id: str) -> None: + async def unregister(self, flow_id: UUID, connection_id: str) -> FlowConnection | None: async with self._lock: room = self._rooms.get(flow_id) if not room: - return - room.pop(connection_id, None) + return None + conn = room.pop(connection_id, None) if not room: self._rooms.pop(flow_id, None) self._remote_rosters.pop(flow_id, None) + self._selections.pop(flow_id, None) + return conn @overload - def local_users(self, flow_id: UUID, *, serialize: Literal[False] = False) -> list[CollaborationPresenceUser]: ... + def local_users( + self, + flow_id: UUID, + *, + serialize: Literal[False] = False, + as_dict: Literal[False] = False, + ) -> list[CollaborationPresenceUser]: ... @overload - def local_users(self, flow_id: UUID, *, serialize: Literal[True]) -> list[dict[str, Any]]: ... + def local_users( + self, + flow_id: UUID, + *, + serialize: Literal[False] = False, + as_dict: Literal[True] = True, + ) -> dict[UUID, CollaborationPresenceUser]: ... + + @overload + def local_users( + self, + flow_id: UUID, + *, + serialize: Literal[True], + as_dict: Literal[False] = False, + ) -> list[dict[str, Any]]: ... + + @overload + def local_users( + self, + flow_id: UUID, + *, + serialize: Literal[True], + as_dict: Literal[True] = True, + ) -> dict[str, dict[str, Any]]: ... def local_users( self, flow_id: UUID, *, serialize: bool = False, - ) -> list[CollaborationPresenceUser] | list[dict[str, Any]]: - """Return unique local users, optionally as JSON-ready payloads.""" + as_dict: bool = False, + ) -> ( + list[CollaborationPresenceUser] + | list[dict[str, Any]] + | dict[UUID, CollaborationPresenceUser] + | dict[str, dict[str, Any]] + ): + """Return unique local users, optionally as a keyed map or JSON-ready payloads.""" room = self._rooms.get(flow_id, {}) if serialize: seen: dict[str, dict[str, Any]] = {} @@ -126,53 +172,85 @@ class CollaborationManager: "username": conn.username, "profile_image": conn.profile_image, } - return list(seen.values()) + return seen if as_dict else list(seen.values()) - seen: dict[str, CollaborationPresenceUser] = {} + seen: dict[UUID, CollaborationPresenceUser] = {} for conn in room.values(): - key = str(conn.user_id) + key = conn.user_id if key not in seen: seen[key] = CollaborationPresenceUser( - user_id=key, + user_id=conn.user_id, username=conn.username, profile_image=conn.profile_image, ) - return list(seen.values()) + return seen if as_dict else list(seen.values()) @overload - def all_users(self, flow_id: UUID, *, serialize: Literal[False] = False) -> list[CollaborationPresenceUser]: ... + def all_users( + self, + flow_id: UUID, + *, + serialize: Literal[False] = False, + as_dict: Literal[False] = False, + ) -> list[CollaborationPresenceUser]: ... @overload - def all_users(self, flow_id: UUID, *, serialize: Literal[True]) -> list[dict[str, Any]]: ... + def all_users( + self, + flow_id: UUID, + *, + serialize: Literal[False] = False, + as_dict: Literal[True] = True, + ) -> dict[UUID, CollaborationPresenceUser]: ... + + @overload + def all_users( + self, + flow_id: UUID, + *, + serialize: Literal[True], + as_dict: Literal[False] = False, + ) -> list[dict[str, Any]]: ... + + @overload + def all_users( + self, + flow_id: UUID, + *, + serialize: Literal[True], + as_dict: Literal[True] = True, + ) -> dict[str, dict[str, Any]]: ... def all_users( self, flow_id: UUID, *, serialize: bool = False, - ) -> list[CollaborationPresenceUser] | list[dict[str, Any]]: - """Return unique local and remote users, optionally as JSON-ready payloads.""" + as_dict: bool = False, + ) -> ( + list[CollaborationPresenceUser] + | list[dict[str, Any]] + | dict[UUID, CollaborationPresenceUser] + | dict[str, dict[str, Any]] + ): + """Return unique local and remote users, optionally as a keyed map or JSON-ready payloads.""" now = time.time() if serialize: - seen: dict[str, dict[str, Any]] = {} - for user in self.local_users(flow_id, serialize=True): - seen[user["user_id"]] = user + seen = self.local_users(flow_id, serialize=True, as_dict=True) for roster in self._remote_rosters.get(flow_id, {}).values(): if now - roster.published_at > PRESENCE_ROSTER_TTL_SECONDS: continue for user in roster.users: - seen[user.user_id] = user.model_dump(mode="json") - return list(seen.values()) + seen[str(user.user_id)] = user.model_dump(mode="json") + return seen if as_dict else list(seen.values()) - seen: dict[str, CollaborationPresenceUser] = {} - for user in self.local_users(flow_id): - seen[user.user_id] = user + seen = self.local_users(flow_id, as_dict=True) for roster in self._remote_rosters.get(flow_id, {}).values(): if now - roster.published_at > PRESENCE_ROSTER_TTL_SECONDS: continue for user in roster.users: seen[user.user_id] = user - return list(seen.values()) + return seen if as_dict else list(seen.values()) def presence_payload(self, flow_id: UUID) -> dict[str, Any]: return { @@ -181,11 +259,68 @@ class CollaborationManager: "users": self.local_users(flow_id, serialize=True), } - def presence_message(self, flow_id: UUID) -> dict[str, Any]: - return { - "type": "presence.updated", - "users": self.all_users(flow_id, serialize=True), - } + def presence_snapshot_message(self, flow_id: UUID) -> dict[str, Any]: + return CollaborationPresenceSnapshotMessage(users=self.all_users(flow_id)).model_dump(mode="json") + + def presence_joined_message( + self, + *, + user_id: UUID, + username: str, + profile_image: str | None, + ) -> dict[str, Any]: + user = CollaborationPresenceUser(user_id=user_id, username=username, profile_image=profile_image) + return CollaborationPresenceJoinedMessage(user=user).model_dump(mode="json") + + 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, flow_id: UUID) -> dict[str, Any]: + selections = [ + CollaborationUserSelection(user_id=user_id, selected=selected) + for user_id, selected in self._selections.get(flow_id, {}).items() + ] + return CollaborationSelectionSnapshotMessage(selections=selections).model_dump(mode="json") + + def selection_updated_message( + self, + user_id: UUID, + selected: CollaborationSelectionTarget | None, + ) -> dict[str, Any]: + return CollaborationSelectionUpdatedMessage( + user_id=user_id, + selected=selected, + ).model_dump(mode="json") + + def set_user_selection( + self, + flow_id: UUID, + user_id: UUID, + selected: CollaborationSelectionTarget | None, + ) -> dict[str, Any]: + flow_selections = self._selections[flow_id] + if selected is None: + flow_selections.pop(user_id, None) + else: + flow_selections[user_id] = selected + return self.selection_updated_message(user_id, selected) + + def clear_user_selection(self, flow_id: UUID, user_id: UUID) -> dict[str, Any] | None: + flow_selections = self._selections.get(flow_id) + if not flow_selections or user_id not in flow_selections: + return None + flow_selections.pop(user_id, None) + return self.selection_updated_message(user_id, None) + + def presence_visibility_diff( + self, + before: dict[UUID, CollaborationPresenceUser], + after: dict[UUID, CollaborationPresenceUser], + ) -> tuple[list[CollaborationPresenceUser], list[UUID]]: + joined_ids = after.keys() - before.keys() + left_ids = before.keys() - after.keys() + joined_users = [after[user_id] for user_id in sorted(joined_ids, key=str)] + return joined_users, sorted(left_ids, key=str) def apply_remote_presence(self, flow_id: UUID, presence: CollaborationPresenceEventPayload) -> None: self._remote_rosters[flow_id][presence.worker_id] = _RemotePresenceRoster( @@ -244,13 +379,26 @@ class CollaborationManager: ) return - if isinstance(backplane_event, CollaborationPresenceUpdatedBackplaneEvent): + if isinstance(backplane_event, CollaborationPresenceRosterBackplaneEvent): payload = backplane_event.payload if payload.worker_id == WORKER_ID: return + before = self.all_users(flow_id, as_dict=True) self.apply_remote_presence(flow_id, payload) - await self.broadcast_json(flow_id, self.presence_message(flow_id)) + after = self.all_users(flow_id, as_dict=True) + joined_users, left_user_ids = self.presence_visibility_diff(before, after) + for user in joined_users: + await self.broadcast_json( + flow_id, + self.presence_joined_message( + user_id=user.user_id, + username=user.username, + profile_image=user.profile_image, + ), + ) + for user_id in left_user_ids: + await self.broadcast_json(flow_id, self.presence_left_message(user_id)) def _prune_fanned_revisions(self) -> None: now = time.time() 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 4fddc736b2..b43d11cbe7 100644 --- a/src/backend/base/langflow/api/v1/schemas/flow_collaboration.py +++ b/src/backend/base/langflow/api/v1/schemas/flow_collaboration.py @@ -11,7 +11,7 @@ from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StrictInt class CollaborationPresenceUser(BaseModel): - user_id: str + user_id: UUID username: str profile_image: str | None = None @@ -25,7 +25,6 @@ class CollaborationSessionReadyMessage(BaseModel): connection_id: str flow_id: UUID current_revision: int - users: list[CollaborationPresenceUser] class CollaborationSessionErrorMessage(BaseModel): @@ -60,11 +59,47 @@ class CollaborationOperationRejectedMessage(BaseModel): current_revision: int | None = None -class CollaborationPresenceUpdatedMessage(BaseModel): - type: Literal["presence.updated"] = "presence.updated" +class CollaborationPresenceSnapshotMessage(BaseModel): + type: Literal["presence.snapshot"] = "presence.snapshot" users: list[CollaborationPresenceUser] +class CollaborationPresenceJoinedMessage(BaseModel): + type: Literal["presence.joined"] = "presence.joined" + user: CollaborationPresenceUser + + +class CollaborationPresenceLeftMessage(BaseModel): + type: Literal["presence.left"] = "presence.left" + user_id: UUID + + +class CollaborationSelectionTarget(BaseModel): + kind: Literal["node", "edge"] + id: str + + +class CollaborationUserSelection(BaseModel): + user_id: UUID + selected: CollaborationSelectionTarget | None = None + + +class CollaborationSelectionUpdateMessage(BaseModel): + type: Literal["selection.update"] + 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 + selected: CollaborationSelectionTarget | None = None + + class CollaborationOperationBroadcastMessage(BaseModel): """Accepted operation broadcast to peers (not echoed to origin).""" @@ -107,14 +142,14 @@ class CollaborationOperationAcceptedBackplaneEvent(BaseModel): payload: CollaborationOperationAcceptedEventPayload -class CollaborationPresenceUpdatedBackplaneEvent(BaseModel): - """Typed presence event consumed from the collaboration backplane.""" +class CollaborationPresenceRosterBackplaneEvent(BaseModel): + """Typed worker-local presence roster event consumed from the collaboration backplane.""" - type: Literal["presence.updated"] + type: Literal["presence.roster"] payload: CollaborationPresenceEventPayload -CollaborationBackplaneEvent = CollaborationOperationAcceptedBackplaneEvent | CollaborationPresenceUpdatedBackplaneEvent +CollaborationBackplaneEvent = CollaborationOperationAcceptedBackplaneEvent | CollaborationPresenceRosterBackplaneEvent class UnsupportedCollaborationBackplaneEventError(ValueError): @@ -126,8 +161,8 @@ def parse_collaboration_backplane_event(event_type: str, payload: dict[str, Any] event = {"type": event_type, "payload": payload} if event_type == "operation.accepted": return CollaborationOperationAcceptedBackplaneEvent.model_validate(event) - if event_type == "presence.updated": - return CollaborationPresenceUpdatedBackplaneEvent.model_validate(event) + if event_type == "presence.roster": + return CollaborationPresenceRosterBackplaneEvent.model_validate(event) msg = f"Unsupported collaboration backplane event type: {event_type!r}" raise UnsupportedCollaborationBackplaneEventError(msg) @@ -137,4 +172,6 @@ class CollaborationUnknownMessageError(BaseModel): detail: str = "Unknown message type" -CollaborationClientMessage = CollaborationSessionStartMessage | CollaborationOperationSubmitMessage +CollaborationClientMessage = ( + CollaborationSessionStartMessage | CollaborationOperationSubmitMessage | CollaborationSelectionUpdateMessage +) 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 55c6af2a7b..a529158508 100644 --- a/src/backend/tests/unit/api/v1/test_flow_collaboration.py +++ b/src/backend/tests/unit/api/v1/test_flow_collaboration.py @@ -99,6 +99,13 @@ def _receive_message_type(ws, expected_type: str) -> dict: return message +def _receive_session_bootstrap(ws) -> tuple[dict, 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 + + def _close_websocket_cleanly(ws) -> None: ws.close() ws.portal.call(anyio.sleep, 0.05) @@ -252,12 +259,16 @@ async def test_collab_session_ready(client: AsyncClient, logged_in_headers): def _assert(ws) -> None: ws.send_json({"type": "session.start"}) - ready = ws.receive_json() + ready, presence, selection = _receive_session_bootstrap(ws) assert ready["type"] == "session.ready" assert ready["current_revision"] == 0 assert ready["flow_id"] == str(flow_id) - assert len(ready["users"]) == 1 - assert ready["users"][0]["username"] == "activeuser" + assert "users" not in ready + 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) @@ -269,7 +280,7 @@ async def test_operation_submit_accepted_increments_revision(client: AsyncClient def _submit(ws) -> None: ws.send_json({"type": "session.start"}) - ready = ws.receive_json() + ready, _, _ = _receive_session_bootstrap(ws) updated = copy.deepcopy(NODE_A) updated["position"] = {"x": 50, "y": 50} ws.send_json( @@ -348,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 = ws.receive_json() + ready, _, _ = _receive_session_bootstrap(ws) ws.send_json( { "type": "operation.submit", @@ -381,7 +392,7 @@ async def test_delete_nodes_removes_incident_edges(client: AsyncClient, logged_i def _submit(ws) -> None: ws.send_json({"type": "session.start"}) - ws.receive_json() + _receive_session_bootstrap(ws) ws.send_json( { "type": "operation.submit", @@ -411,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 = ws.receive_json() + ready, _, _ = _receive_session_bootstrap(ws) ws.send_json( { "type": "operation.submit", @@ -444,9 +455,9 @@ 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 = ws_a.receive_json() + ready_a, _, _ = _receive_session_bootstrap(ws_a) ws_b.send_json({"type": "session.start"}) - ws_b.receive_json() + _receive_session_bootstrap(ws_b) updated = copy.deepcopy(NODE_B) updated["position"] = {"x": 200, "y": 0} @@ -476,11 +487,67 @@ 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"}) - ready_a = ws_a.receive_json() + _, 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"}) - ready_b = ws_b.receive_json() - assert len(ready_a["users"]) == 1 - assert ready_a["users"][0]["username"] == "activeuser" - assert len(ready_b["users"]) == 1 + _, presence_b, _ = _receive_session_bootstrap(ws_b) + assert len(presence_b["users"]) == 1 await _run_dual_websocket_test(app, flow_id, token, _peers) + + +async def test_selection_update_broadcasts_to_peer(client: AsyncClient, logged_in_headers): + token = _access_token(logged_in_headers) + flow_id = await _create_collab_flow(client, logged_in_headers) + app = client._transport.app + + def _peers(ws_a, ws_b) -> None: + ws_a.send_json({"type": "session.start"}) + _receive_session_bootstrap(ws_a) + ws_b.send_json({"type": "session.start"}) + _receive_session_bootstrap(ws_b) + + ws_a.send_json({"type": "selection.update", "selected": {"kind": "node", "id": "a"}}) + updated = _receive_message_type(ws_b, "selection.updated") + assert updated["type"] == "selection.updated" + assert updated["selected"] == {"kind": "node", "id": "a"} + assert updated["user_id"] is not None + + await _run_dual_websocket_test(app, flow_id, token, _peers) + + +async def test_selection_update_supports_null(client: AsyncClient, logged_in_headers): + token = _access_token(logged_in_headers) + flow_id = await _create_collab_flow(client, logged_in_headers) + app = client._transport.app + + def _peers(ws_a, ws_b) -> None: + ws_a.send_json({"type": "session.start"}) + _receive_session_bootstrap(ws_a) + ws_b.send_json({"type": "session.start"}) + _receive_session_bootstrap(ws_b) + + ws_a.send_json({"type": "selection.update", "selected": {"kind": "edge", "id": "e-ab"}}) + _receive_message_type(ws_b, "selection.updated") + ws_a.send_json({"type": "selection.update", "selected": None}) + cleared = _receive_message_type(ws_b, "selection.updated") + assert cleared["selected"] is None + + await _run_dual_websocket_test(app, flow_id, token, _peers) + + +async def test_selection_update_rejects_malformed_payload(client: AsyncClient, logged_in_headers): + token = _access_token(logged_in_headers) + flow_id = await _create_collab_flow(client, logged_in_headers) + app = client._transport.app + + def _assert(ws) -> None: + ws.send_json({"type": "session.start"}) + _receive_session_bootstrap(ws) + ws.send_json({"type": "selection.update", "selected": {"kind": "viewport", "id": "x"}}) + error = ws.receive_json() + assert error["type"] == "message.error" + + await _run_websocket_test(app, flow_id, token, _assert) 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 6a91a0715d..b31c9fda0a 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 @@ -8,7 +8,10 @@ from uuid import uuid4 import pytest from langflow.api.v1 import collaboration_manager as collaboration_manager_module from langflow.api.v1.collaboration_manager import CollaborationManager -from langflow.api.v1.schemas.flow_collaboration import CollaborationPresenceEventPayload +from langflow.api.v1.schemas.flow_collaboration import ( + CollaborationPresenceEventPayload, + CollaborationSelectionTarget, +) from langflow.services.collaboration_events.schemas import CollaborationEvent @@ -60,6 +63,26 @@ async def test_presence_dedupes_same_user_multiple_tabs(manager, flow_id, user_a assert users[0].username == "alice" +@pytest.mark.asyncio +async def test_presence_getters_can_return_keyed_rosters(manager, flow_id, user_a, user_b): + await _register(manager, flow_id, user_a, "alice") + manager.apply_remote_presence( + flow_id, + CollaborationPresenceEventPayload( + worker_id="remote-worker", + published_at=9_999_999_999.0, + users=[{"user_id": str(user_b), "username": "bob", "profile_image": None}], + ), + ) + + local_users = manager.local_users(flow_id, as_dict=True) + all_users = manager.all_users(flow_id, as_dict=True) + + assert set(local_users) == {user_a} + assert set(all_users) == {user_a, user_b} + assert all_users[user_b].username == "bob" + + @pytest.mark.asyncio async def test_presence_payload_publishes_only_local_roster(manager, flow_id, user_a, user_b): await _register(manager, flow_id, user_a, "alice") @@ -182,6 +205,100 @@ async def test_handle_backplane_event_ignores_unknown_type(manager, flow_id, use peer_ws.send_json.assert_not_called() +@pytest.mark.asyncio +async def test_presence_visibility_diff_detects_join_and_leave(manager, flow_id, user_a, user_b): + await _register(manager, flow_id, user_a, "alice") + before = manager.all_users(flow_id, as_dict=True) + await _register(manager, flow_id, user_b, "bob") + after = manager.all_users(flow_id, as_dict=True) + joined_users, left_user_ids = manager.presence_visibility_diff(before, after) + assert [user.username for user in joined_users] == ["bob"] + assert left_user_ids == [] + + +@pytest.mark.asyncio +async def test_presence_snapshot_and_selection_snapshot_messages(manager, flow_id, user_a): + await _register(manager, flow_id, user_a, "alice") + manager.set_user_selection( + flow_id, + user_a, + CollaborationSelectionTarget(kind="node", id="node-1"), + ) + + presence = manager.presence_snapshot_message(flow_id) + selection = manager.selection_snapshot_message(flow_id) + + 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"}}] + + +@pytest.mark.asyncio +async def test_clear_user_selection_returns_updated_message(manager, flow_id, user_a): + await _register(manager, flow_id, user_a, "alice") + manager.set_user_selection( + flow_id, + user_a, + CollaborationSelectionTarget(kind="edge", id="edge-1"), + ) + + cleared = manager.clear_user_selection(flow_id, user_a) + + assert cleared == { + "type": "selection.updated", + "user_id": str(user_a), + "selected": None, + } + assert manager.selection_snapshot_message(flow_id)["selections"] == [] + + +@pytest.mark.asyncio +async def test_handle_backplane_presence_broadcasts_incremental_joined(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 + peer_ws.send_json.reset_mock() + + event = CollaborationEvent( + id="evt-remote-presence", + flow_id=flow_id, + created_at=1.0, + type="presence.roster", + payload={ + "worker_id": "remote-worker", + "published_at": 9_999_999_999.0, + "users": [{"user_id": str(user_a), "username": "alice", "profile_image": None}], + }, + ) + + await manager.handle_backplane_event(event) + + peer_ws.send_json.assert_not_called() + + event = CollaborationEvent( + id="evt-remote-presence-2", + flow_id=flow_id, + created_at=2.0, + type="presence.roster", + payload={ + "worker_id": "remote-worker", + "published_at": 9_999_999_999.0, + "users": [ + {"user_id": str(user_a), "username": "alice", "profile_image": None}, + {"user_id": str(uuid4()), "username": "carol", "profile_image": None}, + ], + }, + ) + + await manager.handle_backplane_event(event) + + peer_ws.send_json.assert_called_once() + payload = peer_ws.send_json.call_args.args[0] + assert payload["type"] == "presence.joined" + assert payload["user"]["username"] == "carol" + + @pytest.mark.asyncio async def test_handle_backplane_presence_ignores_malformed_payload(manager, flow_id, user_a, user_b): await _register(manager, flow_id, user_a, "alice") @@ -192,7 +309,7 @@ async def test_handle_backplane_presence_ignores_malformed_payload(manager, flow id="evt-bad-presence", flow_id=flow_id, created_at=1.0, - type="presence.updated", + type="presence.roster", payload={ "worker_id": "remote-worker", "published_at": "not-a-timestamp", @@ -215,7 +332,7 @@ async def test_handle_backplane_presence_ignores_current_worker_payload(manager, id="evt-own-presence", flow_id=flow_id, created_at=1.0, - type="presence.updated", + type="presence.roster", payload={ "worker_id": collaboration_manager_module.WORKER_ID, "published_at": 9_999_999_999.0, diff --git a/src/backend/tests/unit/services/test_collaboration_events.py b/src/backend/tests/unit/services/test_collaboration_events.py index afef9fadf5..3bd904df89 100644 --- a/src/backend/tests/unit/services/test_collaboration_events.py +++ b/src/backend/tests/unit/services/test_collaboration_events.py @@ -38,13 +38,13 @@ def test_factory_returns_sqlite_implementation(): def test_publish_and_poll_by_flow_id(svc: SQLiteCollaborationEventService, flow_id: UUID): svc.publish(flow_id, "operation.accepted", {"revision": 1, "forward_ops": []}) - svc.publish(flow_id, "presence.updated", {"users": []}) + svc.publish(flow_id, "presence.roster", {"users": []}) events, cursor = svc.poll(flow_id) assert len(events) == 2 assert events[0].type == "operation.accepted" assert events[0].payload["revision"] == 1 - assert events[1].type == "presence.updated" + assert events[1].type == "presence.roster" assert cursor.event_id == events[1].id assert cursor.created_at == events[1].created_at @@ -131,7 +131,7 @@ def test_cross_worker_visibility(tmp_path, flow_id: UUID): assert len(events) == 1 assert events[0].payload["revision"] == 1 - worker_b.publish(flow_id, "presence.updated", {"users": []}) + worker_b.publish(flow_id, "presence.roster", {"users": []}) events, _ = worker_a.poll(flow_id) assert len(events) == 2 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 new file mode 100644 index 0000000000..51c8cc59a0 --- /dev/null +++ b/src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts @@ -0,0 +1,76 @@ +import { + applyPresenceJoined, + applyPresenceLeft, + applyPresenceSnapshot, + applySelectionSnapshot, + applySelectionUpdated, +} from "@/hooks/flows/flow-collaboration-state"; + +describe("flow-collaboration-state", () => { + it("should replace the roster on presence.snapshot", () => { + expect( + applyPresenceSnapshot( + [{ user_id: "old", username: "old-user" }], + [{ user_id: "new", username: "new-user" }], + ), + ).toEqual([{ user_id: "new", username: "new-user" }]); + }); + + it("should add or update users on presence.joined", () => { + expect( + applyPresenceJoined([], { + user_id: "user-1", + username: "ana", + }), + ).toEqual([{ user_id: "user-1", username: "ana" }]); + + expect( + applyPresenceJoined([{ user_id: "user-1", username: "old" }], { + user_id: "user-1", + username: "ana", + profile_image: "Space/046-rocket.svg", + }), + ).toEqual([ + { + user_id: "user-1", + username: "ana", + profile_image: "Space/046-rocket.svg", + }, + ]); + }); + + it("should remove users on presence.left", () => { + expect( + applyPresenceLeft( + [ + { user_id: "user-1", username: "ana" }, + { user_id: "user-2", username: "bob" }, + ], + "user-1", + ), + ).toEqual([{ user_id: "user-2", username: "bob" }]); + }); + + it("should replace selections on selection.snapshot", () => { + expect( + applySelectionSnapshot( + [{ user_id: "user-1", selected: { kind: "node", id: "n1" } }], + [{ user_id: "user-2", selected: { kind: "edge", id: "e1" } }], + ), + ).toEqual([{ user_id: "user-2", selected: { kind: "edge", id: "e1" } }]); + }); + + it("should upsert or clear selections on selection.updated", () => { + expect( + applySelectionUpdated([], "user-1", { kind: "node", id: "n1" }), + ).toEqual([{ user_id: "user-1", selected: { kind: "node", id: "n1" } }]); + + expect( + applySelectionUpdated( + [{ user_id: "user-1", selected: { kind: "node", id: "n1" } }], + "user-1", + null, + ), + ).toEqual([]); + }); +}); diff --git a/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration-editing.test.ts b/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration-editing.test.ts index 44ddece82e..f82651b209 100644 --- a/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration-editing.test.ts +++ b/src/frontend/src/hooks/flows/__tests__/use-flow-collaboration-editing.test.ts @@ -37,8 +37,10 @@ jest.mock("../use-flow-collaboration", () => ({ connectionId: "conn-1", currentRevision: 0, users: [], + selections: [], isReady: true, submitOperations: mockSubmitOperations, + sendSelectionUpdate: jest.fn(), disconnect: mockDisconnect, }; }), 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 3527e2f4ad..4697051176 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 @@ -128,6 +128,9 @@ async function connectSession(currentRevision = 0) { connection_id: "conn-1", flow_id: "flow-1", current_revision: currentRevision, + }); + socket.triggerMessage({ + type: "presence.snapshot", users: [ { user_id: "user-1", @@ -136,6 +139,10 @@ async function connectSession(currentRevision = 0) { }, ], }); + socket.triggerMessage({ + type: "selection.snapshot", + selections: [], + }); }); return socket; @@ -325,21 +332,76 @@ describe("useFlowCollaboration", () => { expect(result.current.currentRevision).toBe(12); }); - it("should update presence roster on presence.updated", async () => { + it("should maintain presence roster from snapshot and incremental events", async () => { const { result } = await mountHook({ flowId: "flow-1" }); await connectSession(0); await act(async () => { latestSocket().triggerMessage({ - type: "presence.updated", - users: [ - { user_id: "user-1", username: "ana" }, - { user_id: "user-2", username: "bob", profile_image: null }, - ], + type: "presence.joined", + user: { user_id: "user-2", username: "bob", profile_image: null }, }); }); - expect(result.current.users).toHaveLength(2); + expect(result.current.users).toEqual([ + { + user_id: "user-1", + username: "ana", + profile_image: "Space/046-rocket.svg", + }, + { user_id: "user-2", username: "bob", profile_image: null }, + ]); + + await act(async () => { + latestSocket().triggerMessage({ + type: "presence.left", + user_id: "user-2", + }); + }); + + expect(result.current.users).toHaveLength(1); + expect(result.current.users[0]?.user_id).toBe("user-1"); + }); + + it("should maintain selection state from snapshot and incremental updates", async () => { + const { result } = await mountHook({ flowId: "flow-1" }); + await connectSession(0); + + await act(async () => { + latestSocket().triggerMessage({ + type: "selection.updated", + user_id: "user-2", + selected: { kind: "node", id: "node-1" }, + }); + }); + + expect(result.current.selections).toEqual([ + { user_id: "user-2", selected: { kind: "node", id: "node-1" } }, + ]); + + await act(async () => { + latestSocket().triggerMessage({ + type: "selection.updated", + user_id: "user-2", + selected: null, + }); + }); + + expect(result.current.selections).toEqual([]); + }); + + it("should send selection.update when sendSelectionUpdate is called", async () => { + const { result } = await mountHook({ flowId: "flow-1" }); + const socket = await connectSession(0); + + await act(async () => { + result.current.sendSelectionUpdate({ kind: "edge", id: "edge-1" }); + }); + + expect(JSON.parse(socket.sent.at(-1)!)).toEqual({ + type: "selection.update", + selected: { kind: "edge", id: "edge-1" }, + }); }); it("should request reload when the socket closes unexpectedly", async () => { diff --git a/src/frontend/src/hooks/flows/flow-collaboration-state.ts b/src/frontend/src/hooks/flows/flow-collaboration-state.ts new file mode 100644 index 0000000000..d90b52bb91 --- /dev/null +++ b/src/frontend/src/hooks/flows/flow-collaboration-state.ts @@ -0,0 +1,59 @@ +import type { + CollaborationPresenceUser, + CollaborationSelectionTarget, + CollaborationUserSelection, +} from "@/types/flow-collaboration"; + +export function applyPresenceSnapshot( + _currentUsers: CollaborationPresenceUser[], + users: CollaborationPresenceUser[], +): CollaborationPresenceUser[] { + return users.map((user) => ({ ...user })); +} + +export function applyPresenceJoined( + currentUsers: CollaborationPresenceUser[], + user: CollaborationPresenceUser, +): CollaborationPresenceUser[] { + const existingIndex = currentUsers.findIndex( + (entry) => entry.user_id === user.user_id, + ); + if (existingIndex === -1) { + return [...currentUsers, { ...user }]; + } + + const nextUsers = [...currentUsers]; + nextUsers[existingIndex] = { ...user }; + return nextUsers; +} + +export function applyPresenceLeft( + currentUsers: CollaborationPresenceUser[], + userId: string, +): CollaborationPresenceUser[] { + return currentUsers.filter((user) => user.user_id !== userId); +} + +export function applySelectionSnapshot( + _currentSelections: CollaborationUserSelection[], + selections: CollaborationUserSelection[], +): CollaborationUserSelection[] { + return selections.map((selection) => ({ + user_id: selection.user_id, + selected: selection.selected ? { ...selection.selected } : null, + })); +} + +export function applySelectionUpdated( + currentSelections: CollaborationUserSelection[], + userId: string, + selected: CollaborationSelectionTarget | null, +): CollaborationUserSelection[] { + const withoutUser = currentSelections.filter( + (entry) => entry.user_id !== userId, + ); + if (selected === null) { + return withoutUser; + } + return [...withoutUser, { user_id: userId, selected: { ...selected } }]; +} diff --git a/src/frontend/src/hooks/flows/use-flow-collaboration.ts b/src/frontend/src/hooks/flows/use-flow-collaboration.ts index 7ac1054b39..cd44b675e9 100644 --- a/src/frontend/src/hooks/flows/use-flow-collaboration.ts +++ b/src/frontend/src/hooks/flows/use-flow-collaboration.ts @@ -1,4 +1,11 @@ import { useCallback, useEffect, useRef, useState } from "react"; +import { + applyPresenceJoined, + applyPresenceLeft, + applyPresenceSnapshot, + applySelectionSnapshot, + applySelectionUpdated, +} from "@/hooks/flows/flow-collaboration-state"; import { buildFlowCollaborationWebSocketUrl } from "@/hooks/flows/flow-collaboration-url"; import type { CollaborationConnectionStatus, @@ -7,8 +14,10 @@ import type { CollaborationPresenceUser, CollaborationReloadDetail, CollaborationReloadReason, + CollaborationSelectionTarget, CollaborationServerMessage, CollaborationSessionErrorMessage, + CollaborationUserSelection, } from "@/types/flow-collaboration"; import type { FlowOperation } from "@/types/flow-operations"; @@ -34,11 +43,13 @@ export type UseFlowCollaborationReturn = { connectionId: string | null; currentRevision: number | null; users: CollaborationPresenceUser[]; + selections: CollaborationUserSelection[]; isReady: boolean; submitOperations: ( operations: FlowOperation[], options?: { requestId?: string }, ) => Promise; + sendSelectionUpdate: (selected: CollaborationSelectionTarget | null) => void; disconnect: () => void; }; @@ -68,6 +79,9 @@ export function useFlowCollaboration({ const [connectionId, setConnectionId] = useState(null); const [currentRevision, setCurrentRevision] = useState(null); const [users, setUsers] = useState([]); + const [selections, setSelections] = useState( + [], + ); const wsRef = useRef(null); const mountedRef = useRef(true); @@ -132,7 +146,6 @@ export function useFlowCollaboration({ case "session.ready": { setConnectionId(message.connection_id); setRevision(message.current_revision); - setUsers(message.users); setStatus("ready"); return; } @@ -186,8 +199,43 @@ export function useFlowCollaboration({ handleRemoteBroadcast(message); return; } - case "presence.updated": { - setUsers(message.users); + case "presence.snapshot": { + setUsers((currentUsers) => + applyPresenceSnapshot(currentUsers, message.users), + ); + return; + } + case "presence.joined": { + setUsers((currentUsers) => + applyPresenceJoined(currentUsers, message.user), + ); + return; + } + case "presence.left": { + setUsers((currentUsers) => + applyPresenceLeft(currentUsers, message.user_id), + ); + setSelections((currentSelections) => + currentSelections.filter( + (entry) => entry.user_id !== message.user_id, + ), + ); + return; + } + case "selection.snapshot": { + setSelections((currentSelections) => + applySelectionSnapshot(currentSelections, message.selections), + ); + return; + } + case "selection.updated": { + setSelections((currentSelections) => + applySelectionUpdated( + currentSelections, + message.user_id, + message.selected, + ), + ); return; } default: @@ -210,6 +258,7 @@ export function useFlowCollaboration({ currentRevisionRef.current = null; setCurrentRevision(null); setUsers([]); + setSelections([]); setStatus("idle"); }, [rejectAllPending]); @@ -225,6 +274,7 @@ export function useFlowCollaboration({ currentRevisionRef.current = null; setCurrentRevision(null); setUsers([]); + setSelections([]); const ws = new WebSocket(buildFlowCollaborationWebSocketUrl(flowId)); wsRef.current = ws; @@ -298,6 +348,22 @@ export function useFlowCollaboration({ [], ); + const sendSelectionUpdate = useCallback( + (selected: CollaborationSelectionTarget | null) => { + const ws = wsRef.current; + if (!ws || ws.readyState !== WebSocket.OPEN) { + return; + } + ws.send( + JSON.stringify({ + type: "selection.update", + selected, + }), + ); + }, + [], + ); + useEffect(() => { mountedRef.current = true; @@ -326,9 +392,11 @@ export function useFlowCollaboration({ connectionId, currentRevision, users, + selections, isReady: status === "ready" && connectionId !== null && currentRevision !== null, submitOperations, + sendSelectionUpdate, disconnect, }; } diff --git a/src/frontend/src/types/flow-collaboration/index.ts b/src/frontend/src/types/flow-collaboration/index.ts index 0e91c73fa2..3ecfc04aca 100644 --- a/src/frontend/src/types/flow-collaboration/index.ts +++ b/src/frontend/src/types/flow-collaboration/index.ts @@ -19,7 +19,6 @@ export type CollaborationSessionReadyMessage = { connection_id: string; flow_id: string; current_revision: number; - users: CollaborationPresenceUser[]; }; export type CollaborationSessionErrorMessage = { @@ -57,11 +56,47 @@ export type CollaborationOperationBroadcastMessage = { created_at: string; }; -export type CollaborationPresenceUpdatedMessage = { - type: "presence.updated"; +export type CollaborationPresenceSnapshotMessage = { + type: "presence.snapshot"; users: CollaborationPresenceUser[]; }; +export type CollaborationPresenceJoinedMessage = { + type: "presence.joined"; + user: CollaborationPresenceUser; +}; + +export type CollaborationPresenceLeftMessage = { + type: "presence.left"; + user_id: string; +}; + +export type CollaborationSelectionTarget = { + kind: "node" | "edge"; + id: string; +}; + +export type CollaborationUserSelection = { + user_id: string; + selected: CollaborationSelectionTarget | null; +}; + +export type CollaborationSelectionUpdateMessage = { + type: "selection.update"; + selected: CollaborationSelectionTarget | null; +}; + +export type CollaborationSelectionSnapshotMessage = { + type: "selection.snapshot"; + selections: CollaborationUserSelection[]; +}; + +export type CollaborationSelectionUpdatedMessage = { + type: "selection.updated"; + user_id: string; + selected: CollaborationSelectionTarget | null; +}; + export type CollaborationMessageErrorMessage = { type: "message.error"; detail: string; @@ -73,12 +108,17 @@ export type CollaborationServerMessage = | CollaborationOperationAcceptedMessage | CollaborationOperationRejectedMessage | CollaborationOperationBroadcastMessage - | CollaborationPresenceUpdatedMessage + | CollaborationPresenceSnapshotMessage + | CollaborationPresenceJoinedMessage + | CollaborationPresenceLeftMessage + | CollaborationSelectionSnapshotMessage + | CollaborationSelectionUpdatedMessage | CollaborationMessageErrorMessage; export type CollaborationClientMessage = | CollaborationSessionStartMessage - | CollaborationOperationSubmitMessage; + | CollaborationOperationSubmitMessage + | CollaborationSelectionUpdateMessage; export type CollaborationConnectionStatus = | "idle"