diff --git a/src/backend/base/langflow/api/utils/collab/connection.py b/src/backend/base/langflow/api/utils/collab/connection.py index b2d4c7cc39..c778c73d3b 100644 --- a/src/backend/base/langflow/api/utils/collab/connection.py +++ b/src/backend/base/langflow/api/utils/collab/connection.py @@ -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) diff --git a/src/backend/base/langflow/api/v1/collaboration_manager.py b/src/backend/base/langflow/api/v1/collaboration_manager.py index 3413a7c5f9..9344d8536d 100644 --- a/src/backend/base/langflow/api/v1/collaboration_manager.py +++ b/src/backend/base/langflow/api/v1/collaboration_manager.py @@ -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) diff --git a/src/backend/base/langflow/services/collaboration_events/service.py b/src/backend/base/langflow/services/collaboration_events/service.py index 333d345f55..198ca71cb5 100644 --- a/src/backend/base/langflow/services/collaboration_events/service.py +++ b/src/backend/base/langflow/services/collaboration_events/service.py @@ -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.""" diff --git a/src/backend/base/langflow/services/collaboration_events/sqlite.py b/src/backend/base/langflow/services/collaboration_events/sqlite.py index 561427a3e0..0038fef921 100644 --- a/src/backend/base/langflow/services/collaboration_events/sqlite.py +++ b/src/backend/base/langflow/services/collaboration_events/sqlite.py @@ -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 diff --git a/src/backend/tests/unit/api/v1/test_collaboration_heartbeat.py b/src/backend/tests/unit/api/v1/test_collaboration_heartbeat.py index 622adf1e19..c056ebef78 100644 --- a/src/backend/tests/unit/api/v1/test_collaboration_heartbeat.py +++ b/src/backend/tests/unit/api/v1/test_collaboration_heartbeat.py @@ -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"), diff --git a/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py b/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py index abe200f532..0ac065b2ee 100644 --- a/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py +++ b/src/backend/tests/unit/api/v1/test_flow_collaboration_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, 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)}, diff --git a/src/backend/tests/unit/services/test_collaboration_events.py b/src/backend/tests/unit/services/test_collaboration_events.py index f938544c31..de892d9b90 100644 --- a/src/backend/tests/unit/services/test_collaboration_events.py +++ b/src/backend/tests/unit/services/test_collaboration_events.py @@ -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()