INNER CODE UNIT · Python
__init__
SHAdd0WTAka/Zen-Ai-Pentest · api/main.py:1113
def __init__(self):
self.active_connections: dict[str, WebSocket] = {}
self.agent_info: dict[str, dict] = {}
async def connect(self, websocket: WebSocket, agent_id: str, api_key: str = None):
await websocket.accept()
self.active_connections[agent_id] = websocket
self.agent_info[agent_id] = {
"api_key": api_key,
"connected_at": datetime.now(timezone.utc).isoformat(),
"last_seen": datetime.now(timezone.utc).isoformat(),
}
logger.info(f"Agent {sanitize_for_log(agent_id)} connected via WebSocket")
def disconnect(self, agent_id: str):
if agent_id in self.active_connections:
del self.active_connections[agent_id]
if agent_id in self.agent_info: