pengjunjie/langflow
1
0
Fork 1
mirror of https://github.com/langflow-ai/langflow.git synced 2026-07-24 00:39:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(collaboration): add incremental presence and selection protocol

Browse Source 
Replace whole-roster WebSocket presence updates with snapshot plus incremental
join/left events, and add ephemeral selection state with bootstrap snapshots and
selection.updated handling. Rename the worker backplane roster event to
presence.roster and simplify connection-side membership checks.
This commit is contained in:
Hamza Rashid
2026-05-31 19:41:09 +00:00
parent 2312fcc322
commit 5ade4e50ee
12 changed files with 814 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
src/backend/base/langflow/api/utils/collab/connection.py
View File

@ -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)

214
src/backend/base/langflow/api/v1/collaboration_manager.py
View File

@ -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()

59
src/backend/base/langflow/api/v1/schemas/flow_collaboration.py
View File

@ -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
)

95
src/backend/tests/unit/api/v1/test_flow_collaboration.py
View File

@ -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)

123
src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py
View File

@ -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,

6
src/backend/tests/unit/services/test_collaboration_events.py
View File

@ -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

76
src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts Normal file
View File

@ -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([]);
});
});

2
src/frontend/src/hooks/flows/__tests__/use-flow-collaboration-editing.test.ts
View File

@ -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,
};
}),

76
src/frontend/src/hooks/flows/__tests__/use-flow-collaboration.test.ts
View File

@ -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 () => {

59
src/frontend/src/hooks/flows/flow-collaboration-state.ts Normal file
View File

@ -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 } }];
}

74
src/frontend/src/hooks/flows/use-flow-collaboration.ts
View File

@ -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<CollaborationOperationAcceptedMessage>;
sendSelectionUpdate: (selected: CollaborationSelectionTarget | null) => void;
disconnect: () => void;
};
@ -68,6 +79,9 @@ export function useFlowCollaboration({
const [connectionId, setConnectionId] = useState<string | null>(null);
const [currentRevision, setCurrentRevision] = useState<number | null>(null);
const [users, setUsers] = useState<CollaborationPresenceUser[]>([]);
const [selections, setSelections] = useState<CollaborationUserSelection[]>(
[],
);
const wsRef = useRef<WebSocket | null>(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,
};
}

50
src/frontend/src/types/flow-collaboration/index.ts
View File

@ -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"