Добавлена генерация id для подключений и сообщений
This commit is contained in:
parent
9c44cc7800
commit
7d93caa42f
1 changed files with 18 additions and 0 deletions
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue