отдельная сущность AgentChat для разделения чатов. Условный "семафор", гарантирующий, что с чатом может работать только одно подключение

src/api/external.py

@@ -1,3 +1,5 @@
+from typing import Annotated
+
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
 
 from lambda_agent_api.server import (
@@ -8,7 +10,8 @@ from lambda_agent_api.server import (
 )
 from lambda_agent_api.client import ClientMessage, MsgUserMessage
 
-from src.agent import get_agent_service, AgentService
+from src.agent import AgentChat
+from src.api.dependencies import get_chat
 
 
 router = APIRouter()
@@ -17,8 +20,7 @@ router = APIRouter()
 @router.websocket("/agent_ws/{chat_id}/")
 async def websocket_endpoint(
     ws: WebSocket,
-    chat_id: int,
-    agent_service: AgentService = Depends(get_agent_service),
+    chat: Annotated[AgentChat, Depends(get_chat)],
 ):
     await ws.accept()
     await ws.send_text(MsgStatus().model_dump_json())
@@ -27,7 +29,7 @@ async def websocket_endpoint(
         while True:
             raw = await ws.receive_text()
             msg = ClientMessage.validate_json(raw)
-            await process_message(ws, chat_id, msg, agent_service)
+            await process_message(ws, chat, msg)
 
     except WebSocketDisconnect:
         pass
@@ -37,9 +39,9 @@ async def websocket_endpoint(
         )
 
 
-async def process_message(ws: WebSocket, chat_id: int, msg, agent_service: AgentService):
+async def process_message(ws: WebSocket, chat: AgentChat, msg):
     match msg:
         case MsgUserMessage():
-            async for chunk in agent_service.astream(chat_id, msg.text):
+            async for chunk in chat.astream(msg.text):
                 await ws.send_text(chunk.model_dump_json())
             await ws.send_text(MsgEventEnd(tokens_used=0).model_dump_json())