From 7fde19b36ad1e0f5ce18a24e5fcbcdad464c16d4 Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Wed, 10 May 2023 12:17:52 -0300 Subject: [PATCH] fix(chat_manager.py): handle case where client is not in active_connections dictionary and remove duplicate code in finally block --- src/backend/langflow/api/chat_manager.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/api/chat_manager.py b/src/backend/langflow/api/chat_manager.py index 1d5cc2406..2141b405b 100644 --- a/src/backend/langflow/api/chat_manager.py +++ b/src/backend/langflow/api/chat_manager.py @@ -4,6 +4,7 @@ from collections import defaultdict from typing import Dict, List from fastapi import WebSocket, status +from fastapi.websockets import WebSocketState from langflow.api.schemas import ChatMessage, ChatResponse, FileResponse from langflow.cache import cache_manager @@ -179,11 +180,13 @@ class ChatManager: await self.active_connections[client_id].close( code=status.WS_1011_INTERNAL_ERROR, reason=str(e) ) + self.disconnect(client_id) finally: try: - await self.active_connections[client_id].close( - code=1000, reason="Client disconnected" - ) + connection = self.active_connections.get(client_id) + if connection: + await connection.close(code=1000, reason="Client disconnected") + self.disconnect(client_id) except Exception as e: logger.exception(e) self.disconnect(client_id)