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): reconcile presence via periodic snapshots

Browse Source 
Reuse the collaboration poll loop to purge expired SQLite connections every 30s and broadcast presence.snapshot per active flow, so clients recover from missed presence.left events. Batch list_users across flow IDs, key connections by connection_id with non-null selected_at, and fold selection into presence.snapshot on the wire and frontend.
This commit is contained in:
Hamza Rashid
2026-06-01 20:22:11 +00:00
parent 3f30023cc1
commit 0dcd3cd978
13 changed files with 398 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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

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

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

4
src/backend/base/langflow/services/collaboration_events/service.py
View File

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

140
src/backend/base/langflow/services/collaboration_events/sqlite.py
View File

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

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

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

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

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

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

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

24
src/frontend/src/hooks/flows/__tests__/flow-collaboration-state.test.ts
View File

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

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

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

24
src/frontend/src/hooks/flows/flow-collaboration-state.ts
View File

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

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

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

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

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