Accepts connection, authenticates user, verifies agent access, loads models, resolves session & history.
(self)
| 288 | await manager.disconnect(str(self.agent_id), self.websocket) |
| 289 | |
| 290 | async def setup(self) -> bool: |
| 291 | """Accepts connection, authenticates user, verifies agent access, loads models, resolves session & history.""" |
| 292 | # Accept immediately so browser sees onopen without waiting for DB setup |
| 293 | await self.websocket.accept() |
| 294 | |
| 295 | # Authenticate |
| 296 | try: |
| 297 | payload = decode_access_token(self.token) |
| 298 | user_id = uuid.UUID(payload["sub"]) |
| 299 | except Exception: |
| 300 | await self.websocket.send_json({"type": "error", "content": "Authentication failed"}) |
| 301 | await self.websocket.close(code=4001) |
| 302 | return False |
| 303 | |
| 304 | try: |
| 305 | async with async_session() as db: |
| 306 | result = await db.execute(select(User).where(User.id == user_id)) |
| 307 | self.user = result.scalar_one_or_none() |
| 308 | if not self.user: |
| 309 | logger.error("[WS] User not found") |
| 310 | await self.websocket.send_json({"type": "error", "content": "User not found"}) |
| 311 | await self.websocket.close(code=4001) |
| 312 | return False |
| 313 | |
| 314 | logger.info(f"[WS] Checking agent access for {self.agent_id}") |
| 315 | self.agent, _ = await check_agent_access(db, self.user, self.agent_id) |
| 316 | if is_agent_expired(self.agent): |
| 317 | await self.websocket.send_json( |
| 318 | { |
| 319 | "type": "error", |
| 320 | "content": "This Agent has expired and is off duty. Please contact your admin to extend its service.", |
| 321 | } |
| 322 | ) |
| 323 | await self.websocket.close(code=4003) |
| 324 | return False |
| 325 | |
| 326 | self.agent_name = self.agent.name |
| 327 | self.agent_type = self.agent.agent_type or "" |
| 328 | self.role_description = self.agent.role_description or "" |
| 329 | self.welcome_message = self.agent.welcome_message or "" |
| 330 | self.ctx_size = self.agent.context_window_size or 100 |
| 331 | self.user_display_name = (self.user.display_name or "").strip() or "there" |
| 332 | logger.info( |
| 333 | f"[WS] Agent: {self.agent_name}, type: {self.agent_type}, model_id: {self.agent.primary_model_id}, ctx: {self.ctx_size}" |
| 334 | ) |
| 335 | |
| 336 | # Load models |
| 337 | await self._load_models(db) |
| 338 | |
| 339 | # Resolve or create chat session |
| 340 | self.conv_id = await self._resolve_chat_session(db, user_id) |
| 341 | if not self.conv_id: |
| 342 | return False |
| 343 | |
| 344 | # Load history messages |
| 345 | await self._load_history(db) |
| 346 | |
| 347 | except Exception as e: |
no test coverage detected