Update app.py

app.py

@@ -1,23 +1,22 @@
 import asyncio
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
-from fastapi.responses import HTMLResponse
+import uuid
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from typing import List, Dict
+from typing import List, Dict, Any
 import uvicorn
 
 # Initialiser l'application FastAPI
 app = FastAPI()
 
 # Configurer CORS pour autoriser toutes les origines
-# Permet à n'importe quel site web d'appeler votre API
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Autorise toutes les origines
+    allow_origins=["*"],
     allow_credentials=True,
-    allow_methods=["*"],  # Autorise toutes les méthodes (GET, POST, etc.)
-    allow_headers=["*"],  # Autorise tous les en-têtes
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
 # Monter un répertoire statique pour servir le fichier index.html
@@ -28,114 +27,103 @@ class MockRequest(BaseModel):
     parameter: str
 
 class ConnectionManager:
-    """Gère les connexions WebSocket actives."""
+    """Gère les connexions WebSocket actives et les réponses en attente."""
     def __init__(self):
         self.active_connections: List[WebSocket] = []
-        # Dictionnaire pour attendre les réponses des clients
+        # Utilise un ID de requête pour lier les requêtes aux réponses
         self.response_futures: Dict[str, asyncio.Future] = {}
 
     async def connect(self, websocket: WebSocket):
-        """Accepte une nouvelle connexion WebSocket."""
         await websocket.accept()
         self.active_connections.append(websocket)
         print(f"Nouvelle connexion WebSocket. Total: {len(self.active_connections)}")
 
     def disconnect(self, websocket: WebSocket):
-        """Ferme une connexion WebSocket."""
+        # Annuler tous les futures en attente pour ce client déconnecté
+        keys_to_remove = [key for key, (ws, _) in self.response_futures.items() if ws == websocket]
+        for key in keys_to_remove:
+            self.response_futures[key][1].cancel()
+            del self.response_futures[key]
+        
         self.active_connections.remove(websocket)
         print(f"Déconnexion WebSocket. Total: {len(self.active_connections)}")
 
-    async def broadcast(self, message: str):
-        """Envoie un message à tous les clients connectés."""
-        # Pour ce cas simple, nous n'envoyons qu'au premier client connecté
-        if self.active_connections:
-            websocket = self.active_connections[0]
-            await websocket.send_text(message)
-            # Créer un Future pour attendre la réponse
-            future = asyncio.get_event_loop().create_future()
-            # Utilise l'identifiant du client comme clé, bien que simple, c'est plus robuste
-            client_id = str(id(websocket))
-            self.response_futures[client_id] = future
-            return future
-        return None
-
-manager = ConnectionManager()
+    async def send_action_and_wait(self, action: str, data: Any):
+        """Envoie une action JSON au client et attend une réponse."""
+        if not self.active_connections:
+            return None
 
+        request_id = str(uuid.uuid4())
+        websocket = self.active_connections[0] # Simplification : on envoie au premier client
 
-@app.get("/", response_class=HTMLResponse)
-async def root():
-    """Serve the main HTML page."""
-    try:
-        with open("static/index.html", "r", encoding="utf-8") as f:
-            return HTMLResponse(content=f.read())
-    except FileNotFoundError:
-        raise HTTPException(status_code=404, detail="index.html not found")
+        future = asyncio.get_event_loop().create_future()
+        self.response_futures[request_id] = future
+        
+        message_to_send = {"request_id": request_id, "action": action, "data": data}
+        await websocket.send_json(message_to_send)
+        
+        return await future
 
+manager = ConnectionManager()
 
-@app.post("/v1/mock")
-async def mock_endpoint(payload: MockRequest):
-    """
-    Endpoint API qui prend un string, le transmet via WebSocket,
-    attend une réponse et la retourne.
-    """
+async def handle_api_request(action: str, payload: MockRequest):
+    """Factorise la logique commune aux endpoints API."""
     try:
-        # Récupérer les données JSON du corps de la requête
-        # data = await request.json()
         input_string = payload.parameter
-
-        if input_string is None:
-            return {"error": "Le paramètre 'parameter' est manquant."}
-            
-        print(f"Endpoint /v1/mock appelé avec: '{input_string}'")
+        print(f"Endpoint pour l'action '{action}' appelé avec: '{input_string}'")
 
         if not manager.active_connections:
             return {"error": "Aucun client WebSocket n'est connecté."}
 
-        # Envoyer le message via WebSocket et obtenir un "future" pour la réponse
-        print("Envoi du message au client WebSocket...")
-        response_future = await manager.broadcast(input_string)
-
-        if response_future is None:
-             return {"error": "Échec de la diffusion du message."}
+        print(f"Envoi de l'action '{action}' au client WebSocket...")
+        websocket_response = await manager.send_action_and_wait(action, input_string)
+        
+        if websocket_response is None:
+             return {"error": "Échec de la communication avec le client."}
 
-        try:
-            # Attendre la réponse du client WebSocket avec un timeout de 10 secondes
-            websocket_response = await asyncio.wait_for(response_future, timeout=10.0)
-            print(f"Réponse reçue du WebSocket: '{websocket_response}'")
-            return {"response_from_client": websocket_response}
-
-        except asyncio.TimeoutError:
-            print("Timeout: Aucune réponse du client WebSocket.")
-            return {"error": "Timeout: Le client n'a pas répondu à temps."}
+        print(f"Réponse reçue du WebSocket pour l'action '{action}': '{websocket_response}'")
+        return {"response_from_client": websocket_response}
 
+    except asyncio.CancelledError:
+        print("La tâche de réponse a été annulée (probablement déconnexion du client).")
+        return {"error": "La requête a été annulée car le client s'est déconnecté."}
     except Exception as e:
-        print(f"Erreur dans /v1/mock: {e}")
+        print(f"Erreur dans handle_api_request: {e}")
         return {"error": f"Une erreur interne est survenue: {str(e)}"}
 
+@app.post("/v1/mock")
+async def mock_endpoint(payload: MockRequest):
+    """Demande au client une phrase prédéfinie."""
+    return await handle_api_request("get_sentence", payload)
+
+@app.post("/v1/reverse")
+async def reverse_endpoint(payload: MockRequest):
+    """Demande au client d'inverser une chaîne de caractères."""
+    return await handle_api_request("reverse_string", payload)
+
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    """Gère la communication WebSocket avec le client."""
     await manager.connect(websocket)
     try:
         while True:
-            # Attendre un message du client
-            data = await websocket.receive_text()
-            print(f"Message reçu du client: '{data}'")
-
-            # Trouver le "future" correspondant et y mettre le résultat
-            client_id = str(id(websocket))
-            if client_id in manager.response_futures:
-                manager.response_futures[client_id].set_result(data)
-                del manager.response_futures[client_id] # Nettoyer après utilisation
+            # Attend une réponse JSON du client
+            response_data = await websocket.receive_json()
+            request_id = response_data.get("request_id")
+            response_payload = response_data.get("response")
+            
+            if request_id and request_id in manager.response_futures:
+                # Marque le future comme terminé avec le résultat
+                manager.response_futures[request_id].set_result(response_payload)
+                del manager.response_futures[request_id]
+            else:
+                print(f"Réponse reçue avec un ID inconnu ou manquant: {request_id}")
 
     except WebSocketDisconnect:
         manager.disconnect(websocket)
-        print("Client déconnecté.")
     except Exception as e:
         print(f"Erreur dans le WebSocket: {e}")
         manager.disconnect(websocket)
 
-
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)