From 7d93caa42f4d5249ecbd7dd60b8569959d44bc4c Mon Sep 17 00:00:00 2001
From: collhoun <2904yr@mail.ru>
Date: Fri, 1 May 2026 23:44:29 +0300
Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?=
 =?UTF-8?q?=D0=BD=D0=B0=20=D0=B3=D0=B5=D0=BD=D0=B5=D1=80=D0=B0=D1=86=D0=B8?=
 =?UTF-8?q?=D1=8F=20id=20=D0=B4=D0=BB=D1=8F=20=D0=BF=D0=BE=D0=B4=D0=BA?=
 =?UTF-8?q?=D0=BB=D1=8E=D1=87=D0=B5=D0=BD=D0=B8=D0=B9=20=D0=B8=20=D1=81?=
 =?UTF-8?q?=D0=BE=D0=BE=D0=B1=D1=89=D0=B5=D0=BD=D0=B8=D0=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/api/external.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

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):