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

refactor(collaboration): make collaboration events service async

Browse Source 
Convert the collaboration backplane and presence store to async APIs
backed by aiosqlite so WebSocket handlers and background poll loops no
longer block the event loop on SQLite I/O. Update call sites and unit
tests, and add collaboration lifecycle logging for connection and
presence flows.
This commit is contained in:
Hamza Rashid
2026-06-04 17:24:30 +00:00
parent e6f8486cb1
commit b631799897
7 changed files with 474 additions and 369 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -143,18 +143,25 @@ class FlowCollaborationConnection:
)
except CollaborationConnectionLimitExceededError as exc:
await self.websocket.close(code=status.WS_1013_TRY_AGAIN_LATER, reason=str(exc))
await logger.awarning("Collaboration connection limit exceeded for flow %s: %s", self.flow_id, exc)
raise _CollaborationConnectionClosedError from exc
self._registered_connection_id = connection_id
await logger.ainfo(
"Collaboration session started for flow %s by user %s (connection_id: %s)",
self.flow_id,
self.current_user.id,
connection_id,
)
await start_collaboration_background_tasks()
presence_change = self._event_service.add_connection(
presence_change = await self._event_service.add_connection(
flow_id=self.flow_id,
user_id=self.current_user.id,
connection_id=connection_id,
username=self.current_user.username,
profile_image=self.current_user.profile_image,
)
snapshot = self._event_service.list_users([self.flow_id])[self.flow_id]
snapshot = (await self._event_service.list_users([self.flow_id]))[self.flow_id]
await self.websocket.send_json(
CollaborationSessionReadyMessage(
@ -171,6 +178,7 @@ class FlowCollaborationConnection:
try:
submit = CollaborationOperationSubmitMessage.model_validate(raw)
except ValidationError as exc:
await logger.awarning("Invalid operation submit payload for flow %s: %s", self.flow_id, exc)
await self._send_operation_rejected(
request_id=raw.get("request_id") if isinstance(raw, dict) else None,
status_code=400,
@ -181,6 +189,7 @@ class FlowCollaborationConnection:
try:
accepted = await self._apply_operation(submit)
except FlowOperationApplyError as exc:
await logger.awarning("Operation rejected for flow %s: %s", self.flow_id, exc.detail)
await self._send_operation_rejected(
request_id=submit.request_id,
status_code=exc.status_code,
@ -189,9 +198,10 @@ class FlowCollaborationConnection:
)
return
await logger.adebug("Operation accepted for flow %s (revision: %s)", self.flow_id, accepted.revision)
await self._send_operation_accepted(submit.request_id, accepted)
await self._broadcast_operation(accepted)
self._publish_operation_accepted(accepted)
await self._publish_operation_accepted(accepted)
async def _apply_operation(
self,
@ -273,7 +283,7 @@ class FlowCollaborationConnection:
exclude_connection_id=self.connection_id,
)
def _publish_operation_accepted(self, accepted: AcceptedFlowOperation) -> None:
async def _publish_operation_accepted(self, accepted: AcceptedFlowOperation) -> None:
event_payload = {
"worker_id": WORKER_ID,
"revision": accepted.revision,
@ -282,7 +292,7 @@ class FlowCollaborationConnection:
"forward_ops": accepted.forward_ops,
"created_at": accepted.created_at.isoformat(),
}
self._event_service.publish(self.flow_id, "operation.accepted", event_payload)
await self._event_service.publish(self.flow_id, "operation.accepted", event_payload)
async def _handle_selection_update(self, raw: Any) -> None:
try:
@ -296,7 +306,7 @@ class FlowCollaborationConnection:
return
connection_id = self.connection_id
change = self._event_service.update_connection(
change = await self._event_service.update_connection(
connection_id=connection_id,
selected=update.selected,
)
@ -322,8 +332,13 @@ class FlowCollaborationConnection:
async def _cleanup(self) -> None:
if self._registered_connection_id is not None:
await logger.ainfo(
"Collaboration connection closed for flow %s (connection_id: %s)",
self.flow_id,
self._registered_connection_id,
)
await self.manager.unregister(self._registered_connection_id)
change = self._event_service.remove_connection(
change = await self._event_service.remove_connection(
connection_id=self._registered_connection_id,
)
await self._emit_presence_change(change)

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

@ -149,11 +149,25 @@ class CollaborationManager:
)
raise CollaborationConnectionLimitExceededError(msg)
self.rooms.add(conn)
await logger.adebug(
"Registered collaboration connection %s for flow %s (user_id: %s)",
connection_id,
flow_id,
user_id,
)
return connection_id
async def unregister(self, connection_id: UUID) -> FlowConnection | None:
async with self._lock:
return self.rooms.remove(connection_id)
conn = self.rooms.remove(connection_id)
if conn is not None:
await logger.adebug(
"Unregistered collaboration connection %s for flow %s (user_id: %s)",
connection_id,
conn.flow_id,
conn.user_id,
)
return conn
async def local_connections_snapshot(self) -> list[FlowConnection]:
async with self._lock:
@ -207,7 +221,7 @@ class CollaborationManager:
return
conn.pong_deadline_at = None
event_service.update_connection(connection_id=connection_id)
await event_service.update_connection(connection_id=connection_id)
async def disconnect_connection(
self,
@ -244,10 +258,13 @@ class CollaborationManager:
except Exception as exc: # noqa: BLE001
await logger.adebug("Failed to close collaboration socket after heartbeat timeout: %s", exc)
if removed_connections:
await logger.adebug("Disconnected %d expired collaboration connections", len(removed_connections))
if not store_connection_ids:
return
for presence_change in event_service.remove_connections(store_connection_ids):
for presence_change in await event_service.remove_connections(store_connection_ids):
await self.emit_presence_change(presence_change.flow_id, presence_change.change, event_service)
def presence_snapshot_message(self, snapshot: CollaborationPresenceSnapshot) -> dict[str, Any]:
@ -294,6 +311,9 @@ class CollaborationManager:
) -> None:
async with self._lock:
connections = self.rooms.connections_for_flow(flow_id)
await logger.adebug("Broadcasting message to %d connections for flow %s", len(connections), flow_id)
for conn in connections:
if exclude_connection_id and conn.connection_id == exclude_connection_id:
continue
@ -338,7 +358,7 @@ class CollaborationManager:
message,
exclude_connection_id=exclude_connection_id,
)
event_service.publish(
await event_service.publish(
flow_id,
"presence.joined",
{
@ -354,7 +374,7 @@ class CollaborationManager:
if change.left_user_id:
message = self.presence_left_message(change.left_user_id)
await self.broadcast_json(flow_id, message)
event_service.publish(
await event_service.publish(
flow_id,
"presence.left",
{
@ -393,7 +413,7 @@ class CollaborationManager:
"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)
await event_service.publish(flow_id, "selection.updated", payload)
async def _handle_operation_accepted_backplane_event(
self,
@ -538,11 +558,14 @@ async def _poll_collaboration_events() -> None:
manager = get_collaboration_manager()
event_service: CollaborationEventService = get_collaboration_events_service()
settings_service = get_settings_service()
snapshot_interval = settings_service.settings.collaboration_presence_snapshot_interval
cursors: dict[UUID, CollaborationPollCursor] = {}
next_users_snapshot_at = time.monotonic()
while True:
flow_ids = manager.rooms.active_flow_ids()
# get immutable view of flow room keys via tuple;
# active roooms can change between this initial fetch and aysnc calls
flow_ids: tuple[UUID, ...] = tuple(manager.rooms.active_flow_ids())
if not flow_ids:
await asyncio.sleep(0.5)
continue
@ -550,18 +573,15 @@ async def _poll_collaboration_events() -> None:
now = time.monotonic()
should_fetch_users = now >= next_users_snapshot_at
if should_fetch_users:
snapshot_interval = settings_service.settings.collaboration_presence_snapshot_interval
next_users_snapshot_at = now + snapshot_interval
snapshots = event_service.list_users(flow_ids)
for flow_id, snapshot in snapshots.items():
await manager.broadcast_json(flow_id, manager.presence_snapshot_message(snapshot))
users_per_flow_room = await event_service.list_users(list(flow_ids))
for flow_id, users_in_flow_room in users_per_flow_room.items():
await manager.broadcast_json(flow_id, manager.presence_snapshot_message(users_in_flow_room))
# This loop awaits while processing events, so snapshot flow ids to avoid
# iterating over a live dict view if rooms change mid-iteration.
for flow_id in tuple(flow_ids):
for flow_id in flow_ids:
cursor = cursors.get(flow_id)
events, new_cursor = event_service.poll(flow_id, cursor=cursor)
events, new_cursor = await event_service.poll(flow_id, cursor=cursor)
cursors[flow_id] = new_cursor
for event in events:
await manager.handle_backplane_event(event)

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

@ -31,11 +31,11 @@ class CollaborationEventService(Service, ABC):
name = "collaboration_events_service"
@abstractmethod
def publish(self, flow_id: UUID, event_type: str, payload: dict) -> CollaborationEvent:
async def publish(self, flow_id: UUID, event_type: str, payload: dict) -> CollaborationEvent:
"""Push an event for cross-worker fanout."""
@abstractmethod
def poll(
async def poll(
self,
flow_id: UUID,
*,
@ -45,11 +45,11 @@ class CollaborationEventService(Service, ABC):
"""Return events after ``cursor`` and the updated poll cursor."""
@abstractmethod
def cleanup(self) -> None:
async def cleanup(self) -> None:
"""Force-evict expired events and presence rows. Useful for tests and ops scripts."""
@abstractmethod
def add_connection(
async def add_connection(
self,
*,
flow_id: UUID,
@ -61,7 +61,7 @@ class CollaborationEventService(Service, ABC):
"""Create or replace one active connection row."""
@abstractmethod
def update_connection(
async def update_connection(
self,
*,
connection_id: UUID,
@ -70,7 +70,7 @@ class CollaborationEventService(Service, ABC):
"""Refresh the connection TTL and optionally update its selection columns."""
@abstractmethod
def remove_connection(
async def remove_connection(
self,
*,
connection_id: UUID,
@ -78,9 +78,9 @@ class CollaborationEventService(Service, ABC):
"""Remove one active connection row."""
@abstractmethod
def remove_connections(self, connection_ids: list[UUID]) -> list[CollaborationPresenceChangeEnvelope]:
async def remove_connections(self, connection_ids: list[UUID]) -> list[CollaborationPresenceChangeEnvelope]:
"""Remove active connection rows and return user-visible presence changes keyed by flow id."""
@abstractmethod
def list_users(self, flow_ids: list[UUID]) -> dict[UUID, CollaborationPresenceSnapshot]:
async def list_users(self, flow_ids: list[UUID]) -> dict[UUID, CollaborationPresenceSnapshot]:
"""Return active deduped users plus effective per-user selections keyed by flow id."""

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

@ -1,14 +1,16 @@
from __future__ import annotations
import asyncio
import json
import sqlite3
import tempfile
import threading
import time
import uuid
from pathlib import Path
from uuid import UUID
import aiosqlite
from lfx.log.logger import logger
from langflow.services.collaboration_events.schemas import (
UNSET,
CollaborationEvent,
@ -58,9 +60,8 @@ CREATE INDEX IF NOT EXISTS idx_connections_expires
class SQLiteCollaborationEventService(CollaborationEventService):
"""SQLite WAL-backed collaboration event mailbox keyed by ``flow_id``.
Uses stdlib ``sqlite3`` so multiple uvicorn/gunicorn workers on the same host
can share one cache-dir database file. This is polling-based fanout, not
pub/sub.
Uses ``aiosqlite`` so multiple uvicorn/gunicorn workers on the same host
can share one cache-dir database file.
"""
TTL_SECONDS: float = 120.0
@ -81,20 +82,29 @@ class SQLiteCollaborationEventService(CollaborationEventService):
if presence_ttl_seconds is not None:
self.PRESENCE_TTL_SECONDS = presence_ttl_seconds
self._db_path = cache_dir / "langflow_collaboration.sqlite"
self._conn = sqlite3.connect(
str(self._db_path),
isolation_level=None,
check_same_thread=False,
timeout=5.0,
)
self._lock = threading.Lock()
with self._lock:
self._conn.execute("PRAGMA journal_mode=WAL")
self._conn.execute("PRAGMA synchronous=NORMAL")
self._conn.execute("PRAGMA busy_timeout=5000")
self._conn.executescript(_SCHEMA)
self._conn: aiosqlite.Connection | None = None
self._init_lock = asyncio.Lock()
self._lock = asyncio.Lock()
def publish(self, flow_id: UUID, event_type: str, payload: dict) -> CollaborationEvent:
async def _ensure_conn(self) -> aiosqlite.Connection:
if self._conn is not None:
return self._conn
async with self._init_lock:
if self._conn is not None:
return self._conn
conn = await aiosqlite.connect(
str(self._db_path),
isolation_level=None,
timeout=5.0,
)
await conn.execute("PRAGMA journal_mode=WAL")
await conn.execute("PRAGMA synchronous=NORMAL")
await conn.execute("PRAGMA busy_timeout=5000")
await conn.executescript(_SCHEMA)
self._conn = conn
return conn
async def publish(self, flow_id: UUID, event_type: str, payload: dict) -> CollaborationEvent:
"""Push an event for cross-worker fanout."""
flow_id_key = str(flow_id)
if not event_type:
@ -108,8 +118,9 @@ class SQLiteCollaborationEventService(CollaborationEventService):
event_payload = dict(payload)
payload_json = json.dumps(event_payload)
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
async with self._lock:
conn = await self._ensure_conn()
await conn.execute("BEGIN IMMEDIATE")
try:
now = time.time()
event = CollaborationEvent(
@ -120,8 +131,8 @@ class SQLiteCollaborationEventService(CollaborationEventService):
payload=event_payload,
)
expires_at = now + self.TTL_SECONDS
self._purge_expired_events_locked(now)
self._conn.execute(
await self._purge_expired_events_locked(conn, now)
await conn.execute(
"""
INSERT INTO events
(id, flow_id, created_at, type, payload, expires_at)
@ -129,15 +140,21 @@ class SQLiteCollaborationEventService(CollaborationEventService):
""",
(event_id, flow_id_key, now, event_type, payload_json, expires_at),
)
self._enforce_per_flow_cap_locked(flow_id_key)
self._conn.execute("COMMIT")
except Exception:
self._conn.execute("ROLLBACK")
await self._enforce_per_flow_cap_locked(conn, flow_id_key)
await conn.execute("COMMIT")
await logger.adebug(
"Published collaboration event %s for flow %s (id: %s)", event_type, flow_id_key, event_id
)
except Exception as exc:
await conn.execute("ROLLBACK")
await logger.aerror(
"Failed to publish collaboration event %s for flow %s: %s", event_type, flow_id_key, exc
)
raise
return event
def poll(
async def poll(
self,
flow_id: UUID,
*,
@ -153,33 +170,34 @@ class SQLiteCollaborationEventService(CollaborationEventService):
raise ValueError(msg)
cursor = cursor or CollaborationPollCursor()
with self._lock:
async with self._lock:
conn = await self._ensure_conn()
now = time.time()
self._purge_expired_events_locked(now)
await self._purge_expired_events_locked(conn, now)
with self._lock:
now = time.time()
rows = self._conn.execute(
"""
SELECT id, created_at, type, payload
FROM events
WHERE flow_id = ?
AND expires_at >= ?
AND (
created_at > ?
OR (created_at = ? AND id > ?)
)
ORDER BY created_at ASC, id ASC
LIMIT ?
""",
(
flow_id_key,
now,
cursor.created_at,
cursor.created_at,
cursor.event_id,
poll_limit,
),
rows = await (
await conn.execute(
"""
SELECT id, created_at, type, payload
FROM events
WHERE flow_id = ?
AND expires_at >= ?
AND (
created_at > ?
OR (created_at = ? AND id > ?)
)
ORDER BY created_at ASC, id ASC
LIMIT ?
""",
(
flow_id_key,
now,
cursor.created_at,
cursor.created_at,
cursor.event_id,
poll_limit,
),
)
).fetchall()
events = [
@ -199,13 +217,15 @@ class SQLiteCollaborationEventService(CollaborationEventService):
last = events[-1]
return events, CollaborationPollCursor(created_at=last.created_at, event_id=last.id)
def cleanup(self) -> None:
async def cleanup(self) -> None:
"""Force-evict expired events and presence rows. Useful for tests and ops scripts."""
with self._lock:
async with self._lock:
conn = await self._ensure_conn()
now = time.time()
self._purge_expired_locked(now)
await self._purge_expired_locked(conn, now)
await logger.adebug("Completed collaboration events cleanup")
def add_connection(
async def add_connection(
self,
*,
flow_id: UUID,
@ -222,15 +242,17 @@ class SQLiteCollaborationEventService(CollaborationEventService):
flow_id_key = str(flow_id)
user_id_key = str(user_id)
connection_id_key = str(connection_id)
change: CollaborationPresenceChange | None = None
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
async with self._lock:
conn = await self._ensure_conn()
await conn.execute("BEGIN IMMEDIATE")
try:
now = time.time()
self._purge_expired_events_locked(now)
before_user = self._effective_user_locked(flow_id_key, user_id_key, now)
await self._purge_expired_events_locked(conn, now)
before_user = await self._effective_user_locked(conn, flow_id_key, user_id_key, now)
expires_at = now + self.PRESENCE_TTL_SECONDS
self._conn.execute(
await conn.execute(
"""
INSERT INTO connections (
flow_id, user_id, connection_id, username, profile_image,
@ -256,16 +278,16 @@ class SQLiteCollaborationEventService(CollaborationEventService):
expires_at,
),
)
after_user = self._effective_user_locked(flow_id_key, user_id_key, now)
after_user = await self._effective_user_locked(conn, flow_id_key, user_id_key, now)
change = self._presence_change(user_id, before_user, after_user)
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
except Exception:
self._conn.execute("ROLLBACK")
await conn.execute("ROLLBACK")
raise
return change
def update_connection(
async def update_connection(
self,
*,
connection_id: UUID,
@ -277,33 +299,37 @@ class SQLiteCollaborationEventService(CollaborationEventService):
selection state. Returns `None` when only the TTL changed.
"""
connection_id_key = str(connection_id)
change: CollaborationPresenceChange | None = None
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
async with self._lock:
conn = await self._ensure_conn()
await conn.execute("BEGIN IMMEDIATE")
try:
now = time.time()
self._purge_expired_events_locked(now)
await self._purge_expired_events_locked(conn, now)
row = self._conn.execute(
"""
SELECT flow_id, user_id
FROM connections
WHERE connection_id = ? AND expires_at >= ?
""",
(connection_id_key, now),
row = await (
await conn.execute(
"""
SELECT flow_id, user_id
FROM connections
WHERE connection_id = ? AND expires_at >= ?
""",
(connection_id_key, now),
)
).fetchone()
if row is None:
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
return None
flow_id_key = row[0]
user_id_key = row[1]
user_id = UUID(user_id_key)
before_user = self._effective_user_locked(flow_id_key, user_id_key, now)
before_user = await self._effective_user_locked(conn, flow_id_key, user_id_key, now)
expires_at = now + self.PRESENCE_TTL_SECONDS
if selected is UNSET:
self._conn.execute(
await conn.execute(
"""
UPDATE connections
SET expires_at = ?
@ -320,7 +346,7 @@ class SQLiteCollaborationEventService(CollaborationEventService):
selected_kind = selected.kind
selected_id = selected.id
selected_at = now
self._conn.execute(
await conn.execute(
"""
UPDATE connections
SET expires_at = ?,
@ -338,16 +364,16 @@ class SQLiteCollaborationEventService(CollaborationEventService):
),
)
after_user = self._effective_user_locked(flow_id_key, user_id_key, now)
after_user = await self._effective_user_locked(conn, flow_id_key, user_id_key, now)
change = self._presence_change(user_id, before_user, after_user)
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
except Exception:
self._conn.execute("ROLLBACK")
await conn.execute("ROLLBACK")
raise
return change
def remove_connection(
async def remove_connection(
self,
*,
connection_id: UUID,
@ -358,31 +384,35 @@ class SQLiteCollaborationEventService(CollaborationEventService):
otherwise `None`.
"""
connection_id_key = str(connection_id)
change: CollaborationPresenceChange | None = None
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
async with self._lock:
conn = await self._ensure_conn()
await conn.execute("BEGIN IMMEDIATE")
try:
now = time.time()
self._purge_expired_events_locked(now)
await self._purge_expired_events_locked(conn, now)
row = self._conn.execute(
"""
SELECT flow_id, user_id
FROM connections
WHERE connection_id = ?
""",
(connection_id_key,),
row = await (
await conn.execute(
"""
SELECT flow_id, user_id
FROM connections
WHERE connection_id = ?
""",
(connection_id_key,),
)
).fetchone()
if row is None:
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
return None
flow_id_key = row[0]
user_id_key = row[1]
user_id = UUID(user_id_key)
before_user = self._effective_user_locked(flow_id_key, user_id_key, now)
before_user = await self._effective_user_locked(conn, flow_id_key, user_id_key, now)
self._conn.execute(
await conn.execute(
"""
DELETE FROM connections
WHERE connection_id = ?
@ -390,48 +420,49 @@ class SQLiteCollaborationEventService(CollaborationEventService):
(connection_id_key,),
)
after_user = self._effective_user_locked(flow_id_key, user_id_key, now)
after_user = await self._effective_user_locked(conn, flow_id_key, user_id_key, now)
change = self._presence_change(user_id, before_user, after_user)
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
except Exception:
self._conn.execute("ROLLBACK")
await conn.execute("ROLLBACK")
raise
return change
def remove_connections(self, connection_ids: list[UUID]) -> list[CollaborationPresenceChangeEnvelope]:
async def remove_connections(self, connection_ids: list[UUID]) -> list[CollaborationPresenceChangeEnvelope]:
"""Remove active connection rows in one transaction."""
if not connection_ids:
return []
connection_id_keys = [str(connection_id) for connection_id in connection_ids]
connection_ids_json = json.dumps(connection_id_keys)
changes: list[CollaborationPresenceChangeEnvelope] = []
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
async with self._lock:
conn = await self._ensure_conn()
await conn.execute("BEGIN IMMEDIATE")
try:
now = time.time()
self._purge_expired_events_locked(now)
await self._purge_expired_events_locked(conn, now)
affected_users = self._affected_users_for_connections_locked(connection_ids_json)
affected_users = await self._affected_users_for_connections_locked(conn, connection_ids_json)
if not affected_users:
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
return []
affected_users_json = json.dumps(
[{"flow_id": flow_id_key, "user_id": user_id_key} for flow_id_key, user_id_key in affected_users]
)
before_users = self._effective_users_for_keys_locked(affected_users_json, now)
self._conn.execute(
before_users = await self._effective_users_for_keys_locked(conn, affected_users_json, now)
await conn.execute(
"""
DELETE FROM connections
WHERE connection_id IN (SELECT value FROM json_each(?))
""",
(connection_ids_json,),
)
after_users = self._effective_users_for_keys_locked(affected_users_json, now)
after_users = await self._effective_users_for_keys_locked(conn, affected_users_json, now)
changes: list[CollaborationPresenceChangeEnvelope] = []
for flow_id_key, user_id_key in affected_users:
flow_id = UUID(flow_id_key)
user_id = UUID(user_id_key)
@ -440,14 +471,14 @@ class SQLiteCollaborationEventService(CollaborationEventService):
if change is not None:
changes.append(CollaborationPresenceChangeEnvelope(flow_id=flow_id, change=change))
self._conn.execute("COMMIT")
await conn.execute("COMMIT")
except Exception:
self._conn.execute("ROLLBACK")
await conn.execute("ROLLBACK")
raise
return changes
def list_users(self, flow_ids: list[UUID]) -> dict[UUID, CollaborationPresenceSnapshot]:
async 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"
@ -455,12 +486,14 @@ class SQLiteCollaborationEventService(CollaborationEventService):
flow_id_keys = [str(flow_id) for flow_id in flow_ids]
flow_ids_json = json.dumps(flow_id_keys)
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
async with self._lock:
conn = await self._ensure_conn()
await conn.execute("BEGIN IMMEDIATE")
try:
now = time.time()
self._purge_expired_events_locked(now)
self._conn.execute(
await self._purge_expired_events_locked(conn, now)
await conn.execute(
"""
DELETE FROM connections
WHERE flow_id IN (SELECT value FROM json_each(?))
@ -468,30 +501,32 @@ class SQLiteCollaborationEventService(CollaborationEventService):
""",
(flow_ids_json, now),
)
snapshots = self._presence_snapshots_for_flows_locked(flow_id_keys, flow_ids_json, now)
self._conn.execute("COMMIT")
snapshots = await self._presence_snapshots_for_flows_locked(conn, flow_id_keys, flow_ids_json, now)
await conn.execute("COMMIT")
except Exception:
self._conn.execute("ROLLBACK")
await 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()
async with self._lock:
if self._conn is not None:
await self._conn.close()
self._conn = None
def _purge_expired_locked(self, now: float) -> None:
self._purge_expired_events_locked(now)
self._purge_expired_connections_locked(now)
async def _purge_expired_locked(self, conn: aiosqlite.Connection, now: float) -> None:
await self._purge_expired_events_locked(conn, now)
await self._purge_expired_connections_locked(conn, now)
def _purge_expired_events_locked(self, now: float) -> None:
self._conn.execute("DELETE FROM events WHERE expires_at < ?", (now,))
async def _purge_expired_events_locked(self, conn: aiosqlite.Connection, now: float) -> None:
await 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,))
async def _purge_expired_connections_locked(self, conn: aiosqlite.Connection, now: float) -> None:
await conn.execute("DELETE FROM connections WHERE expires_at < ?", (now,))
def _enforce_per_flow_cap_locked(self, flow_id: str) -> None:
self._conn.execute(
async def _enforce_per_flow_cap_locked(self, conn: aiosqlite.Connection, flow_id: str) -> None:
await conn.execute(
"""
DELETE FROM events
WHERE rowid IN (
@ -504,52 +539,59 @@ class SQLiteCollaborationEventService(CollaborationEventService):
(flow_id, self.MAX_EVENTS_PER_FLOW),
)
def _affected_users_for_connections_locked(self, connection_ids_json: str) -> list[tuple[str, str]]:
rows = self._conn.execute(
"""
WITH requested(connection_id) AS (
SELECT DISTINCT value
FROM json_each(?)
async def _affected_users_for_connections_locked(
self, conn: aiosqlite.Connection, connection_ids_json: str
) -> list[tuple[str, str]]:
rows = await (
await conn.execute(
"""
WITH requested(connection_id) AS (
SELECT DISTINCT value
FROM json_each(?)
)
SELECT DISTINCT c.flow_id, c.user_id
FROM connections AS c
JOIN requested AS r ON r.connection_id = c.connection_id
ORDER BY c.flow_id ASC, c.user_id ASC
""",
(connection_ids_json,),
)
SELECT DISTINCT c.flow_id, c.user_id
FROM connections AS c
JOIN requested AS r ON r.connection_id = c.connection_id
ORDER BY c.flow_id ASC, c.user_id ASC
""",
(connection_ids_json,),
).fetchall()
return [(row[0], row[1]) for row in rows]
def _effective_users_for_keys_locked(
async def _effective_users_for_keys_locked(
self,
conn: aiosqlite.Connection,
user_keys_json: str,
now: float,
) -> dict[tuple[str, str], CollaborationPresenceConnectionUser]:
rows = self._conn.execute(
"""
WITH affected(flow_id, user_id) AS (
SELECT DISTINCT
json_extract(value, '$.flow_id'),
json_extract(value, '$.user_id')
FROM json_each(?)
rows = await (
await conn.execute(
"""
WITH affected(flow_id, user_id) AS (
SELECT DISTINCT
json_extract(value, '$.flow_id'),
json_extract(value, '$.user_id')
FROM json_each(?)
)
SELECT
c.flow_id,
c.user_id,
c.username,
c.profile_image,
c.selected_kind,
c.selected_id,
MAX(c.selected_at)
FROM connections AS c
JOIN affected AS a
ON a.flow_id = c.flow_id
AND a.user_id = c.user_id
WHERE c.expires_at >= ?
GROUP BY c.flow_id, c.user_id
ORDER BY c.flow_id ASC, c.user_id ASC
""",
(user_keys_json, now),
)
SELECT
c.flow_id,
c.user_id,
c.username,
c.profile_image,
c.selected_kind,
c.selected_id,
MAX(c.selected_at)
FROM connections AS c
JOIN affected AS a
ON a.flow_id = c.flow_id
AND a.user_id = c.user_id
WHERE c.expires_at >= ?
GROUP BY c.flow_id, c.user_id
ORDER BY c.flow_id ASC, c.user_id ASC
""",
(user_keys_json, now),
).fetchall()
users: dict[tuple[str, str], CollaborationPresenceConnectionUser] = {}
@ -563,27 +605,30 @@ class SQLiteCollaborationEventService(CollaborationEventService):
)
return users
def _presence_snapshots_for_flows_locked(
async def _presence_snapshots_for_flows_locked(
self,
conn: aiosqlite.Connection,
flow_id_keys: list[str],
flow_ids_json: str,
now: float,
) -> dict[UUID, CollaborationPresenceSnapshot]:
rows = self._conn.execute(
"""
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 flow_id ASC, user_id ASC, selected_at DESC, connection_id ASC
""",
(flow_ids_json, now),
rows = await (
await conn.execute(
"""
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 flow_id ASC, user_id ASC, selected_at DESC, connection_id ASC
""",
(flow_ids_json, now),
)
).fetchall()
snapshots = {UUID(flow_id_key): CollaborationPresenceSnapshot(users=[]) for flow_id_key in flow_id_keys}
@ -608,8 +653,9 @@ class SQLiteCollaborationEventService(CollaborationEventService):
return snapshots
def _effective_user_locked(
async def _effective_user_locked(
self,
conn: aiosqlite.Connection,
flow_id_key: str,
user_id_key: str,
now: float,
@ -620,8 +666,9 @@ class SQLiteCollaborationEventService(CollaborationEventService):
exist, the user is present once and their effective selection comes from
the connection with the latest selection update.
"""
rows = self._conn.execute(
"""
rows = await (
await conn.execute(
"""
SELECT username, profile_image, selected_kind, selected_id, selected_at
FROM connections
WHERE flow_id = ?
@ -629,7 +676,8 @@ class SQLiteCollaborationEventService(CollaborationEventService):
AND expires_at >= ?
ORDER BY selected_at DESC
""",
(flow_id_key, user_id_key, now),
(flow_id_key, user_id_key, now),
)
).fetchall()
if not rows:
return None

45
src/backend/tests/unit/api/v1/test_collaboration_heartbeat.py
View File

@ -79,13 +79,13 @@ async def test_valid_pong_clears_deadline_and_refreshes_presence(manager, flow_i
conn = manager.rooms.get_connection(conn_id)
assert conn is not None
conn.pong_deadline_at = time.time() + 10.0
event_service = Mock()
event_service = AsyncMock()
with patch("langflow.api.v1.collaboration_manager.time.time", return_value=1000.0):
await manager.handle_heartbeat_pong(flow_id, conn_id, event_service)
assert conn.pong_deadline_at is None
event_service.update_connection.assert_called_once_with(connection_id=conn_id)
event_service.update_connection.assert_awaited_once_with(connection_id=conn_id)
@pytest.mark.asyncio
@ -94,12 +94,12 @@ async def test_late_pong_does_not_refresh_presence(manager, flow_id, user_id):
conn = manager.rooms.get_connection(conn_id)
assert conn is not None
conn.pong_deadline_at = 500.0
event_service = Mock()
event_service = AsyncMock()
with patch("langflow.api.v1.collaboration_manager.time.time", return_value=1000.0):
await manager.handle_heartbeat_pong(flow_id, conn_id, event_service)
event_service.update_connection.assert_not_called()
event_service.update_connection.assert_not_awaited()
@pytest.mark.asyncio
@ -108,30 +108,31 @@ async def test_expired_deadline_disconnects_and_emits_presence_left(manager, flo
conn = manager.rooms.get_connection(conn_id)
assert conn is not None
conn.pong_deadline_at = 1.0
event_service = Mock()
event_service.remove_connections.return_value = [
CollaborationPresenceChangeEnvelope(
flow_id=flow_id,
change=CollaborationPresenceChange(left_user_id=user_id),
)
]
event_service = AsyncMock()
event_service.remove_connections = AsyncMock(
return_value=[
CollaborationPresenceChangeEnvelope(
flow_id=flow_id,
change=CollaborationPresenceChange(left_user_id=user_id),
)
]
)
with patch("langflow.api.v1.collaboration_manager.time.time", return_value=1000.0):
await manager.disconnect_expired_heartbeats(event_service)
assert flow_id not in manager.rooms.active_flow_ids()
event_service.remove_connections.assert_called_once_with([conn_id])
event_service.remove_connections.assert_awaited_once_with([conn_id])
ws.close.assert_called_once()
@pytest.mark.asyncio
async def test_disconnect_connection_is_idempotent(manager):
event_service = Mock()
event_service.remove_connection.return_value = None
event_service = AsyncMock()
event_service.remove_connections = AsyncMock(return_value=[])
await manager.disconnect_connection(uuid4(), event_service)
event_service.remove_connection.assert_not_called()
event_service.remove_connections.assert_not_called()
event_service.remove_connections.assert_not_awaited()
@pytest.mark.asyncio
@ -206,10 +207,10 @@ def test_collaboration_settings_reject_invalid_scheduler(monkeypatch, tmp_path):
async def test_heartbeat_pong_handler_checks_access_before_refresh(monkeypatch):
import langflow.api.utils.collab.connection as connection_module
monkeypatch.setattr(connection_module, "get_collaboration_events_service", Mock(return_value=Mock()))
monkeypatch.setattr(connection_module, "get_collaboration_events_service", Mock(return_value=AsyncMock()))
websocket = AsyncMock()
manager = CollaborationManager()
monkeypatch.setattr(connection_module, "get_collaboration_manager", Mock(return_value=manager))
monkeypatch.setattr(connection_module, "get_collaboration_manager", lambda: manager)
current_user = _user_read()
conn_id = await manager.register(
websocket=websocket,
@ -240,10 +241,10 @@ async def test_heartbeat_pong_handler_checks_access_before_refresh(monkeypatch):
async def test_heartbeat_pong_rejects_extra_fields(monkeypatch):
import langflow.api.utils.collab.connection as connection_module
monkeypatch.setattr(connection_module, "get_collaboration_events_service", Mock(return_value=Mock()))
monkeypatch.setattr(connection_module, "get_collaboration_events_service", Mock(return_value=AsyncMock()))
websocket = AsyncMock()
manager = CollaborationManager()
monkeypatch.setattr(connection_module, "get_collaboration_manager", Mock(return_value=manager))
monkeypatch.setattr(connection_module, "get_collaboration_manager", lambda: manager)
connection = connection_module.FlowCollaborationConnection(
websocket=websocket,
flow_id=UUID("00000000-0000-0000-0000-000000000001"),
@ -265,10 +266,10 @@ async def test_heartbeat_pong_rejects_extra_fields(monkeypatch):
async def test_heartbeat_pong_skipped_when_access_denied(monkeypatch):
import langflow.api.utils.collab.connection as connection_module
monkeypatch.setattr(connection_module, "get_collaboration_events_service", Mock(return_value=Mock()))
monkeypatch.setattr(connection_module, "get_collaboration_events_service", Mock(return_value=AsyncMock()))
websocket = AsyncMock()
manager = CollaborationManager()
monkeypatch.setattr(connection_module, "get_collaboration_manager", Mock(return_value=manager))
monkeypatch.setattr(connection_module, "get_collaboration_manager", lambda: manager)
connection = connection_module.FlowCollaborationConnection(
websocket=websocket,
flow_id=UUID("00000000-0000-0000-0000-000000000001"),

6
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, Mock
from unittest.mock import AsyncMock
from uuid import uuid4
import pytest
@ -233,7 +233,7 @@ async def test_emit_presence_change_broadcasts_and_publishes_left(manager, flow_
peer_conn = manager.rooms.get_connection(peer_id)
assert peer_conn is not None
peer_ws = peer_conn.websocket
event_service = Mock()
event_service = AsyncMock()
await manager.emit_presence_change(flow_id, CollaborationPresenceChange(left_user_id=user_a), event_service)
@ -241,7 +241,7 @@ async def test_emit_presence_change_broadcasts_and_publishes_left(manager, flow_
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(
event_service.publish.assert_awaited_once_with(
flow_id,
"presence.left",
{"worker_id": collaboration_manager_module.WORKER_ID, "user_id": str(user_a)},

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

@ -2,7 +2,7 @@
from __future__ import annotations
import time
import asyncio
from uuid import UUID, uuid4
import pytest
@ -29,8 +29,10 @@ def other_flow_id() -> UUID:
@pytest.fixture
def svc(tmp_path):
return SQLiteCollaborationEventService(cache_dir=tmp_path / "cache")
async def svc(tmp_path):
service = SQLiteCollaborationEventService(cache_dir=tmp_path / "cache")
yield service
await service.teardown()
def test_factory_returns_sqlite_implementation():
@ -41,11 +43,15 @@ def test_factory_returns_sqlite_implementation():
assert isinstance(service, SQLiteCollaborationEventService)
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.joined", {"worker_id": "w1", "user": {"user_id": str(uuid4()), "username": "a"}})
async def test_publish_and_poll_by_flow_id(svc: SQLiteCollaborationEventService, flow_id: UUID):
await svc.publish(flow_id, "operation.accepted", {"revision": 1, "forward_ops": []})
await svc.publish(
flow_id,
"presence.joined",
{"worker_id": "w1", "user": {"user_id": str(uuid4()), "username": "a"}},
)
events, cursor = svc.poll(flow_id)
events, cursor = await svc.poll(flow_id)
assert len(events) == 2
assert events[0].type == "operation.accepted"
assert events[0].payload["revision"] == 1
@ -54,12 +60,12 @@ def test_publish_and_poll_by_flow_id(svc: SQLiteCollaborationEventService, flow_
assert cursor.created_at == events[1].created_at
def test_flow_isolation(svc: SQLiteCollaborationEventService, flow_id: UUID, other_flow_id: UUID):
svc.publish(flow_id, "operation.accepted", {"revision": 1})
svc.publish(other_flow_id, "operation.accepted", {"revision": 2})
async def test_flow_isolation(svc: SQLiteCollaborationEventService, flow_id: UUID, other_flow_id: UUID):
await svc.publish(flow_id, "operation.accepted", {"revision": 1})
await svc.publish(other_flow_id, "operation.accepted", {"revision": 2})
events_a, _ = svc.poll(flow_id)
events_b, _ = svc.poll(other_flow_id)
events_a, _ = await svc.poll(flow_id)
events_b, _ = await svc.poll(other_flow_id)
assert len(events_a) == 1
assert events_a[0].payload["revision"] == 1
@ -67,11 +73,11 @@ def test_flow_isolation(svc: SQLiteCollaborationEventService, flow_id: UUID, oth
assert events_b[0].payload["revision"] == 2
def test_poll_cursor_skips_seen_events(svc: SQLiteCollaborationEventService, flow_id: UUID):
first = svc.publish(flow_id, "operation.accepted", {"revision": 1})
svc.publish(flow_id, "operation.accepted", {"revision": 2})
async def test_poll_cursor_skips_seen_events(svc: SQLiteCollaborationEventService, flow_id: UUID):
first = await svc.publish(flow_id, "operation.accepted", {"revision": 1})
await svc.publish(flow_id, "operation.accepted", {"revision": 2})
events, cursor = svc.poll(
events, cursor = await svc.poll(
flow_id,
cursor=CollaborationPollCursor(created_at=first.created_at, event_id=first.id),
)
@ -80,10 +86,10 @@ def test_poll_cursor_skips_seen_events(svc: SQLiteCollaborationEventService, flo
assert cursor.event_id == events[0].id
def test_poll_at_cursor_returns_empty(svc: SQLiteCollaborationEventService, flow_id: UUID):
event = svc.publish(flow_id, "operation.accepted", {"revision": 1})
async def test_poll_at_cursor_returns_empty(svc: SQLiteCollaborationEventService, flow_id: UUID):
event = await svc.publish(flow_id, "operation.accepted", {"revision": 1})
events, cursor = svc.poll(
events, cursor = await svc.poll(
flow_id,
cursor=CollaborationPollCursor(created_at=event.created_at, event_id=event.id),
)
@ -91,78 +97,88 @@ def test_poll_at_cursor_returns_empty(svc: SQLiteCollaborationEventService, flow
assert cursor == CollaborationPollCursor(created_at=event.created_at, event_id=event.id)
def test_ttl_expiry(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_ttl_expiry(svc: SQLiteCollaborationEventService, flow_id: UUID):
svc.TTL_SECONDS = 0.1
svc.publish(flow_id, "operation.accepted", {"revision": 1})
await svc.publish(flow_id, "operation.accepted", {"revision": 1})
time.sleep(0.15)
await asyncio.sleep(0.15)
events, _ = svc.poll(flow_id)
events, _ = await svc.poll(flow_id)
assert events == []
def test_cleanup_removes_expired(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_cleanup_removes_expired(svc: SQLiteCollaborationEventService, flow_id: UUID):
svc.TTL_SECONDS = 0.1
svc.publish(flow_id, "operation.accepted", {"revision": 1})
await svc.publish(flow_id, "operation.accepted", {"revision": 1})
time.sleep(0.15)
svc.cleanup()
await asyncio.sleep(0.15)
await svc.cleanup()
events, _ = svc.poll(flow_id)
events, _ = await svc.poll(flow_id)
assert events == []
def test_per_flow_cap_eviction(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_per_flow_cap_eviction(svc: SQLiteCollaborationEventService, flow_id: UUID):
svc.MAX_EVENTS_PER_FLOW = 3
for revision in range(5):
svc.publish(flow_id, "operation.accepted", {"revision": revision})
await svc.publish(flow_id, "operation.accepted", {"revision": revision})
events, _ = svc.poll(flow_id)
events, _ = await svc.poll(flow_id)
assert len(events) == 3
assert [event.payload["revision"] for event in events] == [2, 3, 4]
def test_cross_worker_visibility(tmp_path, flow_id: UUID):
async def test_cross_worker_visibility(tmp_path, flow_id: UUID):
"""Two service instances sharing a cache_dir see each other's events."""
shared = tmp_path / "shared"
worker_a = SQLiteCollaborationEventService(cache_dir=shared)
worker_b = SQLiteCollaborationEventService(cache_dir=shared)
worker_a.publish(flow_id, "operation.accepted", {"revision": 1})
try:
await worker_a.publish(flow_id, "operation.accepted", {"revision": 1})
events, _ = worker_b.poll(flow_id)
assert len(events) == 1
assert events[0].payload["revision"] == 1
events, _ = await worker_b.poll(flow_id)
assert len(events) == 1
assert events[0].payload["revision"] == 1
worker_b.publish(flow_id, "selection.updated", {"worker_id": "w2", "user_id": str(uuid4()), "selected": None})
await worker_b.publish(
flow_id, "selection.updated", {"worker_id": "w2", "user_id": str(uuid4()), "selected": None}
)
events, _ = worker_a.poll(flow_id)
assert len(events) == 2
events, _ = await worker_a.poll(flow_id)
assert len(events) == 2
finally:
await worker_a.teardown()
await worker_b.teardown()
def test_cross_worker_flow_isolation(tmp_path, flow_id: UUID, other_flow_id: UUID):
async def test_cross_worker_flow_isolation(tmp_path, flow_id: UUID, other_flow_id: UUID):
shared = tmp_path / "shared"
worker_a = SQLiteCollaborationEventService(cache_dir=shared)
worker_b = SQLiteCollaborationEventService(cache_dir=shared)
worker_a.publish(flow_id, "operation.accepted", {"revision": 1})
worker_b.publish(other_flow_id, "operation.accepted", {"revision": 2})
try:
await worker_a.publish(flow_id, "operation.accepted", {"revision": 1})
await worker_b.publish(other_flow_id, "operation.accepted", {"revision": 2})
events_a, _ = worker_b.poll(flow_id)
events_b, _ = worker_a.poll(other_flow_id)
events_a, _ = await worker_b.poll(flow_id)
events_b, _ = await worker_a.poll(other_flow_id)
assert len(events_a) == 1
assert events_a[0].payload["revision"] == 1
assert len(events_b) == 1
assert events_b[0].payload["revision"] == 2
assert len(events_a) == 1
assert events_a[0].payload["revision"] == 1
assert len(events_b) == 1
assert events_b[0].payload["revision"] == 2
finally:
await worker_a.teardown()
await worker_b.teardown()
def test_publish_requires_event_type(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_publish_requires_event_type(svc: SQLiteCollaborationEventService, flow_id: UUID):
with pytest.raises(ValueError, match="event_type"):
svc.publish(flow_id, "", {})
await svc.publish(flow_id, "", {})
def test_presence_change_fields_are_mutually_exclusive():
@ -181,17 +197,17 @@ 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):
async def test_list_users_requires_flow_ids(svc: SQLiteCollaborationEventService):
with pytest.raises(ValueError, match="flow_ids"):
svc.list_users([])
await svc.list_users([])
def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventService, flow_id: UUID):
user_id = uuid4()
conn_a = uuid4()
conn_b = uuid4()
joined = svc.add_connection(
joined = await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_a,
@ -202,10 +218,10 @@ 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])[flow_id]
snapshot = (await svc.list_users([flow_id]))[flow_id]
assert len(snapshot.users) == 1
second = svc.add_connection(
second = await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_b,
@ -213,9 +229,9 @@ def test_add_update_remove_and_list_presence(svc: SQLiteCollaborationEventServic
profile_image=None,
)
assert second is None
assert len(svc.list_users([flow_id])[flow_id].users) == 1
assert len((await svc.list_users([flow_id]))[flow_id].users) == 1
change = svc.update_connection(
change = await svc.update_connection(
connection_id=conn_a,
selected=CollaborationSelectionTarget(kind="node", id="node-1"),
)
@ -223,43 +239,43 @@ 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])[flow_id]
snapshot = (await svc.list_users([flow_id]))[flow_id]
assert snapshot.users[0].selected == CollaborationSelectionTarget(kind="node", id="node-1")
left = svc.remove_connection(connection_id=conn_a)
left = await svc.remove_connection(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])[flow_id].users) == 1
assert len((await svc.list_users([flow_id]))[flow_id].users) == 1
final = svc.remove_connection(connection_id=conn_b)
final = await svc.remove_connection(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])[flow_id].users == []
assert (await svc.list_users([flow_id]))[flow_id].users == []
def test_remove_connections_batches_presence_changes(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_remove_connections_batches_presence_changes(svc: SQLiteCollaborationEventService, flow_id: UUID):
user_id = uuid4()
other_user_id = uuid4()
conn_a = uuid4()
conn_b = uuid4()
conn_c = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_a,
username="alice",
profile_image=None,
)
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_b,
username="alice",
profile_image=None,
)
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=other_user_id,
connection_id=conn_c,
@ -267,7 +283,7 @@ def test_remove_connections_batches_presence_changes(svc: SQLiteCollaborationEve
profile_image=None,
)
changes = svc.remove_connections([conn_a, conn_c])
changes = await svc.remove_connections([conn_a, conn_c])
assert changes == [
CollaborationPresenceChangeEnvelope(
@ -275,25 +291,25 @@ def test_remove_connections_batches_presence_changes(svc: SQLiteCollaborationEve
change=CollaborationPresenceChange(left_user_id=other_user_id),
)
]
assert [user.user_id for user in svc.list_users([flow_id])[flow_id].users] == [user_id]
assert [user.user_id for user in (await svc.list_users([flow_id]))[flow_id].users] == [user_id]
final_changes = svc.remove_connections([conn_b])
final_changes = await svc.remove_connections([conn_b])
assert final_changes == [
CollaborationPresenceChangeEnvelope(
flow_id=flow_id,
change=CollaborationPresenceChange(left_user_id=user_id),
)
]
assert svc.list_users([flow_id])[flow_id].users == []
assert (await svc.list_users([flow_id]))[flow_id].users == []
def test_remove_connections_batches_effective_user_reads(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_remove_connections_batches_effective_user_reads(svc: SQLiteCollaborationEventService, flow_id: UUID):
user_ids = [uuid4() for _ in range(5)]
connection_ids = []
for index, user_id in enumerate(user_ids):
connection_id = uuid4()
connection_ids.append(connection_id)
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=connection_id,
@ -301,12 +317,13 @@ def test_remove_connections_batches_effective_user_reads(svc: SQLiteCollaboratio
profile_image=None,
)
conn = await svc._ensure_conn()
traced_statements: list[str] = []
svc._conn.set_trace_callback(lambda statement: traced_statements.append(statement.strip()))
await conn.set_trace_callback(lambda statement: traced_statements.append(statement.strip()))
try:
changes = svc.remove_connections(connection_ids)
changes = await svc.remove_connections(connection_ids)
finally:
svc._conn.set_trace_callback(None)
await conn.set_trace_callback(None)
connection_reads = [
statement
@ -318,30 +335,30 @@ def test_remove_connections_batches_effective_user_reads(svc: SQLiteCollaboratio
assert {presence_change.change.left_user_id for presence_change in changes} == set(user_ids)
def test_presence_ttl_cleanup(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_presence_ttl_cleanup(svc: SQLiteCollaborationEventService, flow_id: UUID):
svc.PRESENCE_TTL_SECONDS = 0.1
user_id = uuid4()
conn_id = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_id,
username="alice",
profile_image=None,
)
assert len(svc.list_users([flow_id])[flow_id].users) == 1
assert len((await svc.list_users([flow_id]))[flow_id].users) == 1
time.sleep(0.15)
svc.cleanup()
await asyncio.sleep(0.15)
await svc.cleanup()
assert svc.list_users([flow_id])[flow_id].users == []
assert (await svc.list_users([flow_id]))[flow_id].users == []
def test_list_users_removes_expired_connections(svc: SQLiteCollaborationEventService, flow_id: UUID):
async def test_list_users_removes_expired_connections(svc: SQLiteCollaborationEventService, flow_id: UUID):
svc.PRESENCE_TTL_SECONDS = 0.1
user_id = uuid4()
conn_id = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_id,
@ -349,21 +366,21 @@ def test_list_users_removes_expired_connections(svc: SQLiteCollaborationEventSer
profile_image=None,
)
time.sleep(0.15)
await asyncio.sleep(0.15)
snapshots = svc.list_users([flow_id])
snapshots = await svc.list_users([flow_id])
assert snapshots[flow_id].users == []
assert svc.list_users([flow_id])[flow_id].users == []
assert (await svc.list_users([flow_id]))[flow_id].users == []
def test_poll_does_not_consume_expired_presence_before_list_users(
async 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()
conn_id = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_id,
@ -371,15 +388,15 @@ def test_poll_does_not_consume_expired_presence_before_list_users(
profile_image=None,
)
time.sleep(0.15)
svc.publish(flow_id, "operation.accepted", {"revision": 1})
svc.poll(flow_id)
await asyncio.sleep(0.15)
await svc.publish(flow_id, "operation.accepted", {"revision": 1})
await svc.poll(flow_id)
snapshots = svc.list_users([flow_id])
snapshots = await svc.list_users([flow_id])
assert snapshots[flow_id].users == []
def test_list_users_batches_flow_snapshots(
async def test_list_users_batches_flow_snapshots(
svc: SQLiteCollaborationEventService,
flow_id: UUID,
other_flow_id: UUID,
@ -389,7 +406,7 @@ def test_list_users_batches_flow_snapshots(
active_user_id = uuid4()
expired_conn_id = uuid4()
active_conn_id = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=expired_user_id,
connection_id=expired_conn_id,
@ -397,7 +414,7 @@ def test_list_users_batches_flow_snapshots(
profile_image=None,
)
svc.PRESENCE_TTL_SECONDS = 30.0
svc.add_connection(
await svc.add_connection(
flow_id=other_flow_id,
user_id=active_user_id,
connection_id=active_conn_id,
@ -405,63 +422,63 @@ def test_list_users_batches_flow_snapshots(
profile_image=None,
)
time.sleep(0.15)
await asyncio.sleep(0.15)
snapshots = svc.list_users([flow_id, other_flow_id])
snapshots = await 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):
async def test_effective_selection_uses_latest_connection(svc: SQLiteCollaborationEventService, flow_id: UUID):
user_id = uuid4()
conn_a = uuid4()
conn_b = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_a,
username="alice",
profile_image=None,
)
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_b,
username="alice",
profile_image=None,
)
svc.update_connection(
await svc.update_connection(
connection_id=conn_a,
selected=CollaborationSelectionTarget(kind="node", id="old"),
)
time.sleep(0.01)
svc.update_connection(
await asyncio.sleep(0.01)
await svc.update_connection(
connection_id=conn_b,
selected=CollaborationSelectionTarget(kind="edge", id="new"),
)
snapshot = svc.list_users([flow_id])[flow_id]
snapshot = (await 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):
async def test_initial_connection_counts_as_currently_unselected(svc: SQLiteCollaborationEventService, flow_id: UUID):
user_id = uuid4()
conn_a = uuid4()
conn_b = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_a,
username="alice",
profile_image=None,
)
svc.update_connection(
await svc.update_connection(
connection_id=conn_a,
selected=CollaborationSelectionTarget(kind="node", id="old"),
)
time.sleep(0.01)
svc.add_connection(
await asyncio.sleep(0.01)
await svc.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_b,
@ -469,11 +486,11 @@ def test_initial_connection_counts_as_currently_unselected(svc: SQLiteCollaborat
profile_image=None,
)
snapshot = svc.list_users([flow_id])[flow_id]
snapshot = (await svc.list_users([flow_id]))[flow_id]
assert snapshot.users[0].selected is None
def test_new_distinct_user_snapshot_preserves_existing_user_selection(
async def test_new_distinct_user_snapshot_preserves_existing_user_selection(
svc: SQLiteCollaborationEventService,
flow_id: UUID,
):
@ -482,19 +499,19 @@ def test_new_distinct_user_snapshot_preserves_existing_user_selection(
selected_conn_id = uuid4()
joining_conn_id = uuid4()
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=selected_user_id,
connection_id=selected_conn_id,
username="selected-user",
profile_image=None,
)
svc.update_connection(
await svc.update_connection(
connection_id=selected_conn_id,
selected=CollaborationSelectionTarget(kind="node", id="node-1"),
)
svc.add_connection(
await svc.add_connection(
flow_id=flow_id,
user_id=joining_user_id,
connection_id=joining_conn_id,
@ -502,14 +519,14 @@ def test_new_distinct_user_snapshot_preserves_existing_user_selection(
profile_image=None,
)
snapshot = svc.list_users([flow_id])[flow_id]
snapshot = (await svc.list_users([flow_id]))[flow_id]
users_by_id = {user.user_id: user for user in snapshot.users}
assert users_by_id[selected_user_id].selected == CollaborationSelectionTarget(kind="node", id="node-1")
assert users_by_id[joining_user_id].selected is None
def test_cross_worker_presence_visibility(tmp_path, flow_id: UUID):
async def test_cross_worker_presence_visibility(tmp_path, flow_id: UUID):
shared = tmp_path / "shared"
user_id = uuid4()
conn_id = uuid4()
@ -517,14 +534,18 @@ def test_cross_worker_presence_visibility(tmp_path, flow_id: UUID):
worker_a = SQLiteCollaborationEventService(cache_dir=shared)
worker_b = SQLiteCollaborationEventService(cache_dir=shared)
worker_a.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_id,
username="bob",
profile_image=None,
)
try:
await worker_a.add_connection(
flow_id=flow_id,
user_id=user_id,
connection_id=conn_id,
username="bob",
profile_image=None,
)
snapshot = worker_b.list_users([flow_id])[flow_id]
assert len(snapshot.users) == 1
assert snapshot.users[0].username == "bob"
snapshot = (await worker_b.list_users([flow_id]))[flow_id]
assert len(snapshot.users) == 1
assert snapshot.users[0].username == "bob"
finally:
await worker_a.teardown()
await worker_b.teardown()