diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index 83d85b719..18c8a8a74 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -18,6 +18,10 @@ from langflow.services.utils import get_session from langflow.utils.logger import logger from cachetools import LRUCache from sqlmodel import Session +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from langflow.services.chat.manager import ChatManager router = APIRouter(tags=["Chat"]) @@ -33,16 +37,23 @@ async def chat( ): """Websocket endpoint for chat.""" try: + await websocket.accept() user = await get_current_user(token, db) + if not user: + await websocket.close( + code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized" + ) if not user.is_active: - raise HTTPException(status_code=401, detail="Invalid token") - chat_manager = service_manager.get(ServiceType.CHAT_MANAGER) + await websocket.close( + code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized" + ) + + chat_manager: "ChatManager" = service_manager.get(ServiceType.CHAT_MANAGER) if client_id in chat_manager.in_memory_cache: await chat_manager.handle_websocket(client_id, websocket) else: # We accept the connection but close it immediately # if the flow is not built yet - await websocket.accept() message = "Please, build the flow before sending messages" await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=message) except WebSocketException as exc: diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index a8e3e1790..a434fefcb 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -88,6 +88,10 @@ async def get_current_user( ) user_id: UUID = payload.get("sub") # type: ignore token_type: str = payload.get("type") # type: ignore + if expires := payload.get("exp", None): + expires_datetime = datetime.fromtimestamp(expires) + if datetime.now(timezone.utc) > expires_datetime: + raise credentials_exception if user_id is None or token_type: raise credentials_exception