Repository: Significant-Gravitas/Auto-GPT
File: autogpt_platform/backend/backend/server/conn_manager.py (Python)

Code before:
    ExecutionEventType.NODE_EXEC_UPDATE: WSMethod.NODE_EXECUTION_EVENT,
}


class ConnectionManager:
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()
        self.subscriptions: Dict[str, Set[WebSocket]] = {}
        self.user_connections: Dict[str, Set[WebSocket]] = {}

    async def connect_socket(self, websocket: WebSocket, *, user_id: str):
        await websocket.accept()
        self.active_connections.add(websocket)
        if user_id not in self.user_connections:
            self.user_connections[user_id] = set()
        self.user_connections[user_id].add(websocket)

    def disconnect_socket(self, websocket: WebSocket, *, user_id: str):
        self.active_connections.discard(websocket)
        for subscribers in self.subscriptions.values():
            subscribers.discard(websocket)
        user_conns = self.user_connections.get(user_id)
        if user_conns is not None:
            user_conns.discard(websocket)
            if not user_conns:
                del self.user_connections[user_id]

    async def subscribe_graph_exec(
        self, *, user_id: str, graph_exec_id: str, websocket: WebSocket
    ) -> str:
        return await self._subscribe(
            _graph_exec_channel_key(user_id, graph_exec_id=graph_exec_id), websocket
        )

    async def subscribe_graph_execs(
        self, *, user_id: str, graph_id: str, websocket: WebSocket
    ) -> str:
        return await self._subscribe(
            _graph_execs_channel_key(user_id, graph_id=graph_id), websocket
        )

    async def unsubscribe_graph_exec(
        self, *, user_id: str, graph_exec_id: str, websocket: WebSocket
    ) -> str | None:
        return await self._unsubscribe(
            _graph_exec_channel_key(user_id, graph_exec_id=graph_exec_id), websocket
        )

    async def unsubscribe_graph_execs(
        self, *, user_id: str, graph_id: str, websocket: WebSocket
    ) -> str | None:

Review comment (bug):
is this multi threaded? also should we use pop instead of this manual check otherwise it can fail when multiple user_conns check

Code after:
    ExecutionEventType.GRAPH_EXEC_UPDATE: WSMethod.GRAPH_EXECUTION_EVENT,
    ExecutionEventType.NODE_EXEC_UPDATE: WSMethod.NODE_EXECUTION_EVENT,
}


class ConnectionManager:
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()
        self.subscriptions: Dict[str, Set[WebSocket]] = {}
        self.user_connections: Dict[str, Set[WebSocket]] = {}

    async def connect_socket(self, websocket: WebSocket, *, user_id: str):
        await websocket.accept()
        self.active_connections.add(websocket)
        if user_id not in self.user_connections:
            self.user_connections[user_id] = set()
        self.user_connections[user_id].add(websocket)

    def disconnect_socket(self, websocket: WebSocket, *, user_id: str):
        self.active_connections.discard(websocket)
        for subscribers in self.subscriptions.values():
            subscribers.discard(websocket)
        user_conns = self.user_connections.get(user_id)
        if user_conns is not None:
            user_conns.discard(websocket)
            if not user_conns:
                self.user_connections.pop(user_id, None)

    async def subscribe_graph_exec(
        self, *, user_id: str, graph_exec_id: str, websocket: WebSocket
    ) -> str:
        return await self._subscribe(
            _graph_exec_channel_key(user_id, graph_exec_id=graph_exec_id), websocket
        )

    async def subscribe_graph_execs(
        self, *, user_id: str, graph_id: str, websocket: WebSocket
    ) -> str:
        return await self._subscribe(
            _graph_execs_channel_key(user_id, graph_id=graph_id), websocket
        )

    async def unsubscribe_graph_exec(
        self, *, user_id: str, graph_exec_id: str, websocket: WebSocket
    ) -> str | None:
        return await self._unsubscribe(
            _graph_exec_channel_key(user_id, graph_exec_id=graph_exec_id), websocket
        )

    async def unsubscribe_graph_execs(
        self, *, user_id: str, graph_id: str, websocket: WebSocket