I have an app that made as CQRS with diamonds-service handles all client communication via websockets and make only read operation from db and redis cahce and worker handles all the write operations making updates. I wrote all the basic communication, authentication, websocket connection, etc. Now I want in "while True" part of my connection use case to have some sort of command handler that will handle different commands and send corresponding messages to Apache Pulsar using my PulsarManager. Command payload should be as simple as that: {"command": "create_profile", "payload" : {"some": "command_data}}. Here is my connection manager:
import asyncio
import uuid
from typing import Dict

from fastapi import WebSocket
from fastapi.websockets import WebSocketState

from diamond.schemas import ProfileResponseMessage

__all__ = ["ConnectionManager"]


class ConnectionManager:
    def __init__(self) -> None:
        self.active_connections: Dict[uuid.UUID, Dict[str, WebSocket]] = {}
        self.lock: asyncio.Lock = asyncio.Lock()

    async def connect(
        self, auth_user_uuid: uuid.UUID, state: str, websocket: WebSocket
    ):
        async with self.lock:
            current_user_connections = self.active_connections.get(auth_user_uuid, {})
            current_user_connections.update({state: websocket})
            self.active_connections[auth_user_uuid] = current_user_connections

    async def disconnect(self, auth_user_uuid: uuid.UUID, state: str) -> None:
        current_user_connections = self.active_connections.get(auth_user_uuid, None)
        if current_user_connections:
            websocket = current_user_connections.get(state, None)
            if (
                websocket
                and websocket.application_state is not WebSocketState.DISCONNECTED
            ):
                await websocket.close()

            async with self.lock:
                current_user_connections.pop(state, None)
                self.active_connections[auth_user_uuid] = current_user_connections
        if not current_user_connections:
            self.active_connections.pop(auth_user_uuid, None)

    async def send_message(self, auth_user_uuid: uuid.UUID, message: str) -> None:
        current_user_connections = self.active_connections.get(auth_user_uuid, None)
        for websocket in current_user_connections.values():
            if websocket:
                await websocket.send_text(message)

    async def handle_pubsub_message(self, updates: ProfileResponseMessage) -> None:
        await self.send_message(
            auth_user_uuid=updates.auth_user_uuid, message=updates.model_dump_json()
        )


And here is my UseCase class for handling websocket connection:
from fastapi import WebSocket
from fastapi.websockets import WebSocketDisconnect

from usecases import BaseUseCase
from auth.services import TokenServiceInterface
from diamond.services.diamond_service import DiamondServiceInterface

class HandleWebsocketConnectionUseCase(BaseUseCase):
    def __init__(
        self,
        token_service: TokenServiceInterface,
        diamond_service: DiamondServiceInterface,
    ) -> None:
        self._token_service = token_service
        self._diamond_service = diamond_service

    async def execute(self, websocket: WebSocket, connection_state: str) -> None:
        await self._diamond_service.accept_ws_connection(websocket, connection_state)

        token = await self._diamond_service.get_auth_token(
            await websocket.receive_text()
        )

        token_data = await self._token_service.get_access_token_data(token)
        await self._diamond_service.connect_websocket(auth_user_uuid=token_data.auth_user_uuid, state=connection_state, websocket=websocket)

        try:
            while True:
                await websocket.receive_text()
        except WebSocketDisconnect:
            if token_data:
                await self._diamond_service.disconnect_websocket(auth_user_uuid=token_data.auth_user_uuid, state=connection_state)