aditya
9fe3196827
All sims re-run after pipeline fixes (better prompts, real syntax validation, fenced-block extraction fix). Results: 7/7 PASS, 20 files total, all Python files pass compile() syntax check. Includes run_all_sims.py runner script and rag-basic action config.
139 lines
5.4 KiB
Python
139 lines
5.4 KiB
Python
import asyncio
|
|
import websockets
|
|
import json
|
|
import logging
|
|
from typing import Set
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Global set to store connected clients
|
|
connected_clients: Set[websockets.WebSocketServerProtocol] = set()
|
|
|
|
async def register_client(websocket: websockets.WebSocketServerProtocol) -> None:
|
|
"""Register a new client connection."""
|
|
connected_clients.add(websocket)
|
|
logger.info(f"Client {websocket.remote_address} connected. Total clients: {len(connected_clients)}")
|
|
|
|
# Notify all clients about new connection
|
|
if len(connected_clients) > 1:
|
|
notification = {
|
|
"type": "system",
|
|
"message": f"User from {websocket.remote_address[0]} joined the chat",
|
|
"timestamp": asyncio.get_event_loop().time()
|
|
}
|
|
await broadcast_message(json.dumps(notification), exclude=websocket)
|
|
|
|
async def unregister_client(websocket: websockets.WebSocketServerProtocol) -> None:
|
|
"""Unregister a client connection."""
|
|
if websocket in connected_clients:
|
|
connected_clients.remove(websocket)
|
|
logger.info(f"Client {websocket.remote_address} disconnected. Total clients: {len(connected_clients)}")
|
|
|
|
# Notify remaining clients about disconnection
|
|
if connected_clients:
|
|
notification = {
|
|
"type": "system",
|
|
"message": f"User from {websocket.remote_address[0]} left the chat",
|
|
"timestamp": asyncio.get_event_loop().time()
|
|
}
|
|
await broadcast_message(json.dumps(notification))
|
|
|
|
async def broadcast_message(message: str, exclude: websockets.WebSocketServerProtocol = None) -> None:
|
|
"""Broadcast a message to all connected clients except the excluded one."""
|
|
if not connected_clients:
|
|
return
|
|
|
|
# Create a copy of the set to avoid modification during iteration
|
|
clients_to_notify = connected_clients.copy()
|
|
if exclude:
|
|
clients_to_notify.discard(exclude)
|
|
|
|
if clients_to_notify:
|
|
# Send message to all clients concurrently
|
|
await asyncio.gather(
|
|
*[send_safe(client, message) for client in clients_to_notify],
|
|
return_exceptions=True
|
|
)
|
|
|
|
async def send_safe(websocket: websockets.WebSocketServerProtocol, message: str) -> None:
|
|
"""Safely send a message to a client with error handling."""
|
|
try:
|
|
await websocket.send(message)
|
|
except websockets.exceptions.ConnectionClosed:
|
|
logger.warning(f"Failed to send message to {websocket.remote_address}: connection closed")
|
|
await unregister_client(websocket)
|
|
except Exception as e:
|
|
logger.error(f"Error sending message to {websocket.remote_address}: {e}")
|
|
await unregister_client(websocket)
|
|
|
|
async def handle_client(websocket: websockets.WebSocketServerProtocol, path: str) -> None:
|
|
"""Handle a client connection."""
|
|
await register_client(websocket)
|
|
|
|
try:
|
|
async for message in websocket:
|
|
try:
|
|
# Parse the incoming message
|
|
data = json.loads(message)
|
|
|
|
# Create a broadcast message with metadata
|
|
broadcast_data = {
|
|
"type": "chat",
|
|
"message": data.get("message", ""),
|
|
"sender": data.get("sender", "Anonymous"),
|
|
"timestamp": asyncio.get_event_loop().time()
|
|
}
|
|
|
|
logger.info(f"Broadcasting message from {broadcast_data['sender']}: {broadcast_data['message']}")
|
|
|
|
# Broadcast to all connected clients
|
|
await broadcast_message(json.dumps(broadcast_data))
|
|
|
|
except json.JSONDecodeError:
|
|
logger.warning(f"Invalid JSON received from {websocket.remote_address}")
|
|
error_response = {
|
|
"type": "error",
|
|
"message": "Invalid message format",
|
|
"timestamp": asyncio.get_event_loop().time()
|
|
}
|
|
await send_safe(websocket, json.dumps(error_response))
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing message from {websocket.remote_address}: {e}")
|
|
|
|
except websockets.exceptions.ConnectionClosed:
|
|
logger.info(f"Connection closed by {websocket.remote_address}")
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error with {websocket.remote_address}: {e}")
|
|
finally:
|
|
await unregister_client(websocket)
|
|
|
|
async def main():
|
|
"""Start the WebSocket server."""
|
|
host = "localhost"
|
|
port = 8765
|
|
|
|
logger.info(f"Starting WebSocket server on {host}:{port}")
|
|
|
|
try:
|
|
async with websockets.serve(handle_client, host, port):
|
|
logger.info(f"WebSocket server started successfully on ws://{host}:{port}")
|
|
logger.info("Press Ctrl+C to stop the server")
|
|
|
|
# Keep the server running
|
|
await asyncio.Future() # Run forever
|
|
|
|
except OSError as e:
|
|
logger.error(f"Failed to start server: {e}")
|
|
except KeyboardInterrupt:
|
|
logger.info("Server shutdown requested")
|
|
except Exception as e:
|
|
logger.error(f"Unexpected server error: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
logger.info("Server stopped") |