50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
import asyncio
|
|
from uuid import UUID
|
|
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class ConsolePresenceManager:
|
|
def __init__(self) -> None:
|
|
self._connections: dict[UUID, WebSocket] = {}
|
|
self._relay_queues: dict[UUID, asyncio.Queue[str]] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def register(self, user_id: UUID, websocket: WebSocket) -> None:
|
|
async with self._lock:
|
|
self._connections[user_id] = websocket
|
|
|
|
async def unregister(self, user_id: UUID, websocket: WebSocket) -> None:
|
|
async with self._lock:
|
|
if self._connections.get(user_id) is websocket:
|
|
self._connections.pop(user_id, None)
|
|
self._relay_queues.pop(user_id, None)
|
|
|
|
async def is_present(self, user_id: UUID) -> bool:
|
|
async with self._lock:
|
|
return user_id in self._connections
|
|
|
|
async def get_connection(self, user_id: UUID) -> WebSocket | None:
|
|
async with self._lock:
|
|
return self._connections.get(user_id)
|
|
|
|
async def start_relay(self, user_id: UUID) -> asyncio.Queue[str]:
|
|
async with self._lock:
|
|
queue: asyncio.Queue[str] = asyncio.Queue()
|
|
self._relay_queues[user_id] = queue
|
|
return queue
|
|
|
|
async def stop_relay(self, user_id: UUID) -> None:
|
|
async with self._lock:
|
|
self._relay_queues.pop(user_id, None)
|
|
|
|
async def forward_from_console(self, user_id: UUID, text: str) -> bool:
|
|
async with self._lock:
|
|
queue = self._relay_queues.get(user_id)
|
|
if not queue:
|
|
return False
|
|
await queue.put(text)
|
|
return True
|
|
|
|
|
|
console_presence = ConsolePresenceManager()
|