Добавлена генерация id для подключений и сообщений

This commit is contained in:
Ярослав Малинин 2026-05-01 23:44:29 +03:00
parent 9c44cc7800
commit 7d93caa42f

View file

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