WebSockets

Configuración básica de WebSocket

Los WebSockets representan un protocolo de comunicación que permite el intercambio de datos bidireccional entre cliente y servidor en tiempo real. A diferencia de las peticiones HTTP tradicionales, donde el cliente siempre inicia la comunicación, los WebSockets mantienen una conexión persistente que permite a ambas partes enviar mensajes cuando sea necesario.

FastAPI incluye soporte nativo para WebSockets a través del decorador @app.websocket(), lo que facilita enormemente la implementación de funcionalidades en tiempo real como notificaciones instantáneas, actualizaciones de estado o sistemas de chat.

Instalación de dependencias

Para trabajar con WebSockets en FastAPI necesitamos instalar una dependencia adicional que gestione las conexiones WebSocket:

pip install "fastapi[standard]" python-multipart

La dependencia python-multipart es necesaria para manejar correctamente las conexiones WebSocket y los datos que se intercambian a través de ellas.

Estructura básica de un endpoint WebSocket

Un endpoint WebSocket en FastAPI se define de manera similar a los endpoints HTTP, pero utilizando el decorador @app.websocket(). La principal diferencia es que recibe un parámetro WebSocket en lugar de Request:

from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    
    try:
        while True:
            # Recibir mensaje del cliente
            data = await websocket.receive_text()
            
            # Enviar respuesta al cliente
            await websocket.send_text(f"Mensaje recibido: {data}")
            
    except Exception as e:
        print(f"Error en WebSocket: {e}")
    finally:
        await websocket.close()

El método websocket.accept() es obligatorio y debe llamarse antes de poder intercambiar mensajes. Establece la conexión WebSocket con el cliente y confirma que el servidor está listo para recibir datos.

Métodos principales del objeto WebSocket

El objeto WebSocket proporciona varios métodos para gestionar la comunicación:

• await websocket.accept(): Acepta la conexión WebSocket entrante
• await websocket.receive_text(): Recibe un mensaje de texto del cliente
• await websocket.receive_json(): Recibe datos JSON del cliente
• await websocket.send_text(mensaje): Envía un mensaje de texto al cliente
• await websocket.send_json(datos): Envía datos JSON al cliente
• await websocket.close(): Cierra la conexión WebSocket

Ejemplo práctico: Echo Server

Un echo server es el ejemplo más básico para entender el funcionamiento de WebSockets. Simplemente devuelve al cliente el mismo mensaje que recibe:

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

@app.websocket("/echo")
async def echo_websocket(websocket: WebSocket):
    await websocket.accept()
    
    try:
        while True:
            # Esperar mensaje del cliente
            message = await websocket.receive_text()
            
            # Devolver el mismo mensaje
            await websocket.send_text(f"Echo: {message}")
            
    except WebSocketDisconnect:
        print("Cliente desconectado")

La excepción WebSocketDisconnect se produce cuando el cliente cierra la conexión de forma inesperada. Es importante capturarla para evitar errores en el servidor.

Intercambio de datos JSON

Los WebSockets también permiten intercambiar datos estructurados en formato JSON, lo que resulta muy útil para aplicaciones más complejas:

@app.websocket("/json-ws")
async def json_websocket(websocket: WebSocket):
    await websocket.accept()
    
    try:
        while True:
            # Recibir datos JSON
            data = await websocket.receive_json()
            
            # Procesar los datos
            response = {
                "tipo": "respuesta",
                "mensaje": f"Procesado: {data.get('contenido', '')}",
                "timestamp": data.get('timestamp', 0) + 1
            }
            
            # Enviar respuesta JSON
            await websocket.send_json(response)
            
    except WebSocketDisconnect:
        print("Conexión cerrada")

Cliente HTML básico para pruebas

Para probar nuestros endpoints WebSocket, podemos crear un cliente HTML simple que se conecte desde el navegador:

<!DOCTYPE html>
<html>
<head>
    <title>Test WebSocket</title>
</head>
<body>
    <div id="messages"></div>
    <input type="text" id="messageInput" placeholder="Escribe un mensaje">
    <button onclick="sendMessage()">Enviar</button>

    <script>
        const ws = new WebSocket("ws://localhost:8000/echo");
        
        ws.onmessage = function(event) {
            const messages = document.getElementById('messages');
            messages.innerHTML += '<p>' + event.data + '</p>';
        };
        
        function sendMessage() {
            const input = document.getElementById('messageInput');
            ws.send(input.value);
            input.value = '';
        }
    </script>
</body>
</html>

Gestión de errores y desconexiones

Una implementación robusta debe manejar adecuadamente las desconexiones y errores que pueden ocurrir durante la comunicación:

@app.websocket("/robust-ws")
async def robust_websocket(websocket: WebSocket):
    await websocket.accept()
    print("Nueva conexión WebSocket establecida")
    
    try:
        while True:
            message = await websocket.receive_text()
            
            # Validar mensaje antes de procesarlo
            if not message.strip():
                await websocket.send_text("Error: Mensaje vacío")
                continue
                
            await websocket.send_text(f"Procesado: {message}")
            
    except WebSocketDisconnect:
        print("Cliente desconectado normalmente")
    except Exception as e:
        print(f"Error inesperado: {e}")
        try:
            await websocket.close()
        except:
            pass  # La conexión ya puede estar cerrada

Esta configuración básica proporciona la base necesaria para implementar comunicación en tiempo real. El bucle infinito mantiene la conexión activa, mientras que el manejo de excepciones asegura que el servidor continúe funcionando incluso cuando los clientes se desconectan inesperadamente.

Chat simple en tiempo real

Un chat en tiempo real es una de las aplicaciones más comunes de WebSockets, ya que permite que múltiples usuarios intercambien mensajes instantáneamente. Para implementar esta funcionalidad necesitamos gestionar múltiples conexiones simultáneas y distribuir los mensajes entre todos los participantes conectados.

Gestor de conexiones

El primer paso es crear un gestor de conexiones que mantenga un registro de todos los clientes conectados al chat:

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List
import json
from datetime import datetime

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
    
    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
    
    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)
    
    async def broadcast(self, message: str):
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except:
                # Si falla el envío, remover la conexión
                self.active_connections.remove(connection)

manager = ConnectionManager()
app = FastAPI()

La clase ConnectionManager mantiene una lista de conexiones activas y proporciona métodos para conectar, desconectar y enviar mensajes tanto individuales como masivos.

Endpoint del chat

El endpoint principal del chat gestiona las conexiones entrantes y distribuye los mensajes:

@app.websocket("/chat")
async def chat_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    
    # Notificar que un usuario se ha conectado
    await manager.broadcast("Un usuario se ha conectado al chat")
    
    try:
        while True:
            # Recibir mensaje del cliente
            data = await websocket.receive_text()
            
            # Crear mensaje con timestamp
            timestamp = datetime.now().strftime("%H:%M:%S")
            formatted_message = f"[{timestamp}] {data}"
            
            # Enviar mensaje a todos los conectados
            await manager.broadcast(formatted_message)
            
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast("Un usuario ha abandonado el chat")

Chat con nombres de usuario

Para mejorar la experiencia, podemos permitir que los usuarios establezcan un nombre de usuario al conectarse:

class ChatManager:
    def __init__(self):
        self.active_connections: dict = {}  # websocket: username
    
    async def connect(self, websocket: WebSocket, username: str):
        await websocket.accept()
        self.active_connections[websocket] = username
    
    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            del self.active_connections[websocket]
    
    async def broadcast_message(self, message: str, sender_websocket: WebSocket):
        sender_name = self.active_connections.get(sender_websocket, "Anónimo")
        timestamp = datetime.now().strftime("%H:%M:%S")
        
        formatted_message = f"[{timestamp}] {sender_name}: {message}"
        
        for websocket in list(self.active_connections.keys()):
            try:
                await websocket.send_text(formatted_message)
            except:
                self.disconnect(websocket)

chat_manager = ChatManager()

@app.websocket("/chat/{username}")
async def chat_with_username(websocket: WebSocket, username: str):
    await chat_manager.connect(websocket, username)
    
    # Notificar entrada del usuario
    join_message = f"{username} se ha unido al chat"
    await chat_manager.broadcast_system_message(join_message)
    
    try:
        while True:
            message = await websocket.receive_text()
            await chat_manager.broadcast_message(message, websocket)
            
    except WebSocketDisconnect:
        chat_manager.disconnect(websocket)
        leave_message = f"{username} ha abandonado el chat"
        await chat_manager.broadcast_system_message(leave_message)

Mensajes del sistema

Es útil distinguir entre mensajes de usuarios y mensajes del sistema (como notificaciones de entrada y salida):

async def broadcast_system_message(self, message: str):
    timestamp = datetime.now().strftime("%H:%M:%S")
    system_message = f"[{timestamp}] SISTEMA: {message}"
    
    for websocket in list(self.active_connections.keys()):
        try:
            await websocket.send_text(system_message)
        except:
            self.disconnect(websocket)

Cliente HTML para el chat

Un cliente web básico que permita a los usuarios participar en el chat:

<!DOCTYPE html>
<html>
<head>
    <title>Chat en Tiempo Real</title>
    <style>
        #chat-container { 
            max-width: 600px; 
            margin: 20px auto; 
            font-family: Arial, sans-serif; 
        }
        #messages { 
            height: 400px; 
            border: 1px solid #ccc; 
            overflow-y: scroll; 
            padding: 10px; 
            background-color: #f9f9f9; 
        }
        #input-container { 
            display: flex; 
            margin-top: 10px; 
        }
        #messageInput { 
            flex: 1; 
            padding: 10px; 
            border: 1px solid #ccc; 
        }
        #sendButton { 
            padding: 10px 20px; 
            background-color: #007bff; 
            color: white; 
            border: none; 
            cursor: pointer; 
        }
    </style>
</head>
<body>
    <div id="chat-container">
        <h2>Chat en Tiempo Real</h2>
        <div id="messages"></div>
        <div id="input-container">
            <input type="text" id="messageInput" placeholder="Escribe tu mensaje...">
            <button id="sendButton" onclick="sendMessage()">Enviar</button>
        </div>
    </div>

    <script>
        const username = prompt("Ingresa tu nombre de usuario:");
        const ws = new WebSocket(`ws://localhost:8000/chat/${username}`);
        
        ws.onmessage = function(event) {
            const messages = document.getElementById('messages');
            const messageElement = document.createElement('div');
            messageElement.textContent = event.data;
            messages.appendChild(messageElement);
            messages.scrollTop = messages.scrollHeight;
        };
        
        function sendMessage() {
            const input = document.getElementById('messageInput');
            if (input.value.trim()) {
                ws.send(input.value);
                input.value = '';
            }
        }
        
        document.getElementById('messageInput').addEventListener('keypress', function(e) {
            if (e.key === 'Enter') {
                sendMessage();
            }
        });
    </script>
</body>
</html>

Validación y filtrado de mensajes

Para un chat más robusto, podemos agregar validación básica de los mensajes:

def validate_message(message: str) -> bool:
    # Verificar que el mensaje no esté vacío
    if not message.strip():
        return False
    
    # Limitar longitud del mensaje
    if len(message) > 500:
        return False
    
    # Filtrar palabras prohibidas básicas
    prohibited_words = ["spam", "admin"]
    message_lower = message.lower()
    
    for word in prohibited_words:
        if word in message_lower:
            return False
    
    return True

@app.websocket("/chat/{username}")
async def validated_chat(websocket: WebSocket, username: str):
    await chat_manager.connect(websocket, username)
    
    try:
        while True:
            message = await websocket.receive_text()
            
            if validate_message(message):
                await chat_manager.broadcast_message(message, websocket)
            else:
                await websocket.send_text("SISTEMA: Mensaje no válido")
                
    except WebSocketDisconnect:
        chat_manager.disconnect(websocket)

Lista de usuarios conectados

Podemos añadir funcionalidad para mostrar qué usuarios están actualmente conectados:

async def get_connected_users(self) -> List[str]:
    return list(self.active_connections.values())

async def broadcast_user_list(self):
    users = await self.get_connected_users()
    user_list_message = f"SISTEMA: Usuarios conectados: {', '.join(users)}"
    
    for websocket in list(self.active_connections.keys()):
        try:
            await websocket.send_text(user_list_message)
        except:
            self.disconnect(websocket)

Esta implementación básica de chat proporciona comunicación bidireccional en tiempo real entre múltiples usuarios. Los mensajes se distribuyen instantáneamente a todos los participantes conectados, creando una experiencia de chat fluida y responsiva.