Fixed logging

src/api/external.py

@@ -1,10 +1,10 @@
 from typing import Annotated
 
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
+from pydantic_core import ValidationError
 
 from lambda_agent_api.server import (
     MsgStatus,
-    MsgEventTextChunk,
     MsgEventEnd,
     MsgError,
 )
@@ -27,22 +27,27 @@ async def websocket_endpoint(
     # важно использовать именно _ws вариант, чтобы корректно обрабатывались исключения
     chat: Annotated[AgentChat, Depends(get_chat_ws)],
 ):
-    logger.trace(f"WebSocket connection accepted for chat_id: {chat_id}")
+    logger.info(f"WebSocket connection accepted for chat_id: {chat_id}")
     await ws.accept()
     await ws.send_text(MsgStatus().model_dump_json())
 
     try:
         while True:
             raw = await ws.receive_text()
-            logger.trace(f"Received raw message: {raw}")
-            msg = ClientMessage.validate_json(raw)
+            logger.trace(f"Received raw message: {len(raw)} characters for chat_id: {chat_id}")
+            try:
+                msg = ClientMessage.validate_json(raw)
+            except ValidationError as e:
+                logger.warning(f"Invalid JSON received from chat {chat_id}: {e}")
+                await ws.send_text(MsgError(code="BAD_REQUEST", details="Invalid message format").model_dump_json())
+                continue
             await process_message(ws, chat, msg)
 
     except WebSocketDisconnect:
-        logger.trace(f"WebSocket disconnected for chat_id: {chat_id}")
+        logger.info(f"WebSocket disconnected for chat_id: {chat_id}")
         pass
     except Exception as exc:
-        logger.trace(f"Error occurred for chat_id {chat_id}: {exc}")
+        logger.exception("Unexpected error in websocket")
         await ws.send_text(
             MsgError(code="INTERNAL_ERROR", details=str(exc)).model_dump_json()
         )
@@ -51,7 +56,9 @@ async def websocket_endpoint(
 async def process_message(ws: WebSocket, chat: AgentChat, msg):
     match msg:
         case MsgUserMessage():
-            logger.trace(f"Processing user message: {msg.text}")
+            logger.debug(f"Processing user message for chat {chat.chat_id} (text length: {len(msg.text)}, attachments: {len(msg.attachments) if msg.attachments else 0})")
             async for chunk in chat.astream(msg.text, msg.attachments):
-                logger.trace(f"Sending chunk: {chunk}")
+                logger.trace(f"Sending stream chunk to chat {chat.chat_id}: {chunk.__class__.__name__}")
                 await ws.send_text(chunk.model_dump_json())
+            logger.debug(f"Finished processing user message for chat {chat.chat_id}")
+            await ws.send_text(MsgEventEnd(tokens_used=0).model_dump_json())  # TODO: подставить реальное потребление токенов