diff --git a/src/api/external.py b/src/api/external.py index c76e74f..6721bf6 100644 --- a/src/api/external.py +++ b/src/api/external.py @@ -13,6 +13,13 @@ from lambda_agent_api.client import ClientMessage, MsgUserMessage from src.agent import AgentChat from src.api.dependencies import get_chat_ws from src.core.logger import get_logger +from src.core.correlation import ( + generate_connection_id, + generate_message_id, + set_connection_id, + set_message_id, + clear_context, +) router = APIRouter() @@ -27,6 +34,10 @@ async def websocket_endpoint( # важно использовать именно _ws вариант, чтобы корректно обрабатывались исключения chat: Annotated[AgentChat, Depends(get_chat_ws)], ): + # Генерируем уникальный ID для этого подключения + connection_id = generate_connection_id() + set_connection_id(connection_id) + logger.info(f"WebSocket connection accepted for chat_id: {chat_id}") await ws.accept() await ws.send_text(MsgStatus().model_dump_json()) @@ -34,6 +45,11 @@ async def websocket_endpoint( try: while True: raw = await ws.receive_text() + + # Генерируем ID для каждого сообщения + message_id = generate_message_id() + set_message_id(message_id) + logger.trace(f"Received raw message: {len(raw)} characters for chat_id: {chat_id}") try: msg = ClientMessage.validate_json(raw) @@ -51,6 +67,8 @@ async def websocket_endpoint( await ws.send_text( MsgError(code="INTERNAL_ERROR", details=str(exc)).model_dump_json() ) + finally: + clear_context() async def process_message(ws: WebSocket, chat: AgentChat, msg):