Текстовый файл из учебного комплекта, не свидетельство установленного runtime. Вернуться к учебнику · Исходный файл
1 from __future__ import annotations 2 3 import asyncio 4 import json 5 import logging 6 import signal 7 import sys 8 from uuid import uuid4 9 10 from nats.aio.client import Client as NATS 11 from nats.aio.msg import Msg 12 13 from src_agent.agent import AgentRunner 14 from src_agent.repositories.runtime_state import load_runtime_context 15 from src_agent.settings import STACK_SERVICE_NAME 16 from src_agent.ui_events import UiEventPublisher 17 from src_agent.utils.db import Database, fetch_chat_history 18 from src_agent.utils.nats_logger import NatsLogger 19 from src_shared.contracts import ( 20 AgentInboundRequest, 21 AgentInboundResponse, 22 ErrorInfo, 23 HistoryGetRequest, 24 HistoryGetResponse, 25 now_ts_ms, 26 ) 27 28 logger = logging.getLogger(STACK_SERVICE_NAME) 29 30 31 class AgentServer: 32 """ 33 Thin NATS runner: 34 - validates inbound requests 35 - ensures DB session/runtime rows 36 - calls workflow runtime via NATS 37 - persists updated runtime state 38 - publishes UI commands/events 39 """ 40 41 def __init__( 42 self, 43 *, 44 nats_url: str, 45 agent_subject: str, 46 agent_history_subject: str, 47 events_subject: str, 48 workflow_subject: str, 49 workflow_timeout_s: int, 50 workflow_no_responders_retries: int, 51 workflow_no_responders_retry_delay_s: float, 52 db_url: str, 53 user_id: str, 54 runtime_conflict_retries: int, 55 ) -> None: 56 self.nats_url = nats_url 57 self.agent_subject = agent_subject 58 self.agent_history_subject = agent_history_subject 59 self.events_subject = events_subject 60 self.workflow_subject = workflow_subject 61 self.workflow_timeout_s = int(workflow_timeout_s) 62 self.workflow_no_responders_retries = max(0, int(workflow_no_responders_retries)) 63 self.workflow_no_responders_retry_delay_s = max(0.1, float(workflow_no_responders_retry_delay_s)) 64 self.runtime_conflict_retries = int(runtime_conflict_retries) 65 self.db_url = db_url 66 self.user_id = user_id 67 68 self.nc: NATS | None = None 69 self.db: Database | None = None 70 self.nats_logger: NatsLogger | None = None 71 self.runner: AgentRunner | None = None 72 self.ui_publisher: UiEventPublisher | None = None 73 self.stop_event = asyncio.Event() 74 75 async def connect(self) -> None: 76 self.nc = NATS() 77 await self.nc.connect( 78 servers=[self.nats_url], 79 name=STACK_SERVICE_NAME, 80 max_reconnect_attempts=-1, 81 reconnect_time_wait=2, 82 ping_interval=10, 83 ) 84 self.nats_logger = NatsLogger(self.nc, self.events_subject, STACK_SERVICE_NAME) 85 self.ui_publisher = UiEventPublisher(self.nats_logger) 86 87 self.db = Database(db_url=self.db_url) 88 await self.db.connect() 89 self.runner = AgentRunner( 90 nc=self.nc, 91 db=self.db, 92 workflow_subject=self.workflow_subject, 93 user_id=self.user_id, 94 workflow_timeout_s=self.workflow_timeout_s, 95 max_runtime_conflict_retries=self.runtime_conflict_retries, 96 workflow_no_responders_retries=self.workflow_no_responders_retries, 97 workflow_no_responders_retry_delay_s=self.workflow_no_responders_retry_delay_s, 98 ) 99 await self.nats_logger.info(f"{STACK_SERVICE_NAME} connected (nats={self.nats_url})") 100 101 async def subscribe(self) -> None: 102 if not self.nc: 103 raise RuntimeError("NATS is not connected") 104 await self.nc.subscribe( 105 self.agent_subject, 106 queue=f"{STACK_SERVICE_NAME}.q.{self.user_id}", 107 cb=self.handle_request, 108 ) 109 await self.nc.subscribe( 110 self.agent_history_subject, 111 queue=f"{STACK_SERVICE_NAME}.history.q.{self.user_id}", 112 cb=self.handle_history_request, 113 ) 114 if self.nats_logger: 115 await self.nats_logger.info(f"Subscribed to {self.agent_subject} and {self.agent_history_subject}") 116 117 @staticmethod 118 def _agent_exception_response(exc: Exception) -> AgentInboundResponse: 119 return AgentInboundResponse( 120 trace_id=uuid4(), 121 correlation_id=None, 122 request_id=uuid4(), 123 session_id=None, 124 ts_ms=now_ts_ms(), 125 ok=False, 126 status="FAILED", 127 result="", 128 client_handler={}, 129 client_events=[], 130 errors=[ErrorInfo(code="agent_exception", message=str(exc))], 131 ) 132 133 @staticmethod 134 def _history_exception_response(exc: Exception) -> HistoryGetResponse: 135 return HistoryGetResponse( 136 trace_id=uuid4(), 137 correlation_id=None, 138 request_id=uuid4(), 139 session_id=None, 140 ts_ms=now_ts_ms(), 141 ok=False, 142 items=[], 143 runtime_context=None, 144 error=ErrorInfo(code="history_exception", message=str(exc)), 145 ) 146 147 @staticmethod 148 def _response_bytes(payload: AgentInboundResponse | HistoryGetResponse) -> bytes: 149 return payload.model_dump_json().encode("utf-8") 150 151 async def handle_request(self, msg: Msg) -> None: 152 try: 153 raw = json.loads(msg.data.decode("utf-8")) 154 if isinstance(raw, dict) and isinstance(raw.get("text"), str): 155 raw["text"] = raw["text"].strip() 156 req = AgentInboundRequest.model_validate(raw) 157 158 if not self.runner: 159 raise RuntimeError("AgentServer is not initialized") 160 161 rr = await self.runner.run(req) 162 if self.ui_publisher: 163 await self.ui_publisher.publish(req, rr.workflow) 164 165 resp = AgentInboundResponse( 166 trace_id=req.trace_id, 167 correlation_id=req.correlation_id, 168 request_id=req.request_id, 169 session_id=rr.session_id, 170 ts_ms=now_ts_ms(), 171 ok=rr.workflow.status != "FAILED", 172 status=rr.workflow.status, 173 result=rr.workflow.result, 174 client_handler=rr.workflow.client_handler, 175 client_events=rr.workflow.client_events, 176 errors=rr.workflow.errors, 177 ) 178 await msg.respond(self._response_bytes(resp)) 179 180 except Exception as exc: 181 logger.exception("Error processing request") 182 fallback = self._agent_exception_response(exc) 183 try: 184 await msg.respond(self._response_bytes(fallback)) 185 except Exception: 186 logger.exception("Failed to respond with error") 187 if self.nats_logger: 188 await self.nats_logger.error(f"Processing error: {str(exc)}") 189 190 async def handle_history_request(self, msg: Msg) -> None: 191 try: 192 raw = json.loads(msg.data.decode("utf-8")) 193 req = HistoryGetRequest.model_validate(raw) 194 if req.session_id is None: 195 raise ValueError("session_id is required") 196 if self.db is None: 197 raise RuntimeError("AgentServer DB is not initialized") 198 199 items = await fetch_chat_history( 200 self.db, 201 session_id=str(req.session_id), 202 limit=req.limit, 203 ) 204 runtime_context = await load_runtime_context( 205 self.db, 206 session_id=req.session_id, 207 ) 208 resp = HistoryGetResponse( 209 trace_id=req.trace_id, 210 correlation_id=req.correlation_id, 211 request_id=req.request_id, 212 session_id=req.session_id, 213 ts_ms=now_ts_ms(), 214 ok=True, 215 items=items, 216 runtime_context=runtime_context, 217 error=None, 218 ) 219 await msg.respond(self._response_bytes(resp)) 220 except Exception as exc: 221 logger.exception("Error processing history request") 222 fallback = self._history_exception_response(exc) 223 try: 224 await msg.respond(self._response_bytes(fallback)) 225 except Exception: 226 logger.exception("Failed to respond with history error") 227 228 def setup_signal_handlers(self) -> None: 229 loop = asyncio.get_running_loop() 230 for sig in (signal.SIGINT, signal.SIGTERM): 231 loop.add_signal_handler(sig, lambda s=sig: asyncio.create_task(self.shutdown(s))) 232 233 async def shutdown(self, signal_obj=None) -> None: 234 if signal_obj: 235 logger.info("Received exit signal %s", getattr(signal_obj, "name", signal_obj)) 236 self.stop_event.set() 237 238 if self.nc: 239 try: 240 await self.nc.drain() 241 except Exception: 242 logger.exception("Error draining NATS connection") 243 finally: 244 await self.nc.close() 245 logger.info("NATS connection closed") 246 247 if self.db: 248 try: 249 await self.db.close() 250 except Exception: 251 logger.exception("Error closing DB pool") 252 253 async def run(self) -> None: 254 try: 255 await self.connect() 256 await self.subscribe() 257 self.setup_signal_handlers() 258 logger.info("%s is ready", STACK_SERVICE_NAME) 259 await self.stop_event.wait() 260 except Exception as exc: 261 logger.exception("Server error: %s", str(exc)) 262 sys.exit(1) 263 finally: 264 await self.shutdown()