"""Bot Provider Management - Test fix for stale lobby IDs after restart""" import os import time import uuid import secrets import threading import httpx import json import asyncio from typing import Dict, List, Optional from pydantic import ValidationError from logger import logger # Import shared models import sys sys.path.append( os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) ) from shared.models import ( BotProviderModel, BotProviderPublicModel, BotProviderRegisterRequest, BotProviderRegisterResponse, BotProviderListResponse, BotListResponse, BotInfoModel, BotJoinLobbyRequest, BotJoinLobbyResponse, BotLeaveLobbyRequest, BotLeaveLobbyResponse, BotProviderBotsResponse, BotProviderJoinResponse, BotJoinPayload, BotInstanceModel, ) from core.session_manager import SessionManager from core.lobby_manager import LobbyManager class BotProviderConfig: """Configuration class for bot provider management""" # Comma-separated list of allowed provider keys # Format: "key1:name1,key2:name2" or just "key1,key2" (names default to keys) ALLOWED_PROVIDERS = os.getenv("BOT_PROVIDER_KEYS", "") @classmethod def get_allowed_providers(cls) -> Dict[str, str]: """Parse allowed providers from environment variable Returns: dict mapping provider_key -> provider_name """ if not cls.ALLOWED_PROVIDERS.strip(): return {} providers: Dict[str, str] = {} for entry in cls.ALLOWED_PROVIDERS.split(","): entry = entry.strip() if not entry: continue if ":" in entry: key, name = entry.split(":", 1) providers[key.strip()] = name.strip() else: providers[entry] = entry return providers class BotManager: """Manages bot providers and bot lifecycle""" def __init__(self): self.bot_providers: Dict[str, BotProviderModel] = {} self.bot_instances: Dict[ str, BotInstanceModel ] = {} # bot_instance_id -> BotInstanceModel self.lock = threading.RLock() self.bot_providers_file = "bot_providers.json" self.cleanup_task: Optional["asyncio.Task[None]"] = None self._shutdown_event = asyncio.Event() # Check if provider authentication is enabled allowed_providers = BotProviderConfig.get_allowed_providers() if not allowed_providers: logger.warning("Bot provider authentication disabled. Any provider can register.") # Load persisted bot providers self._load_bot_providers() # Note: Don't start cleanup task here - will be started when needed def start_cleanup(self): """Start the cleanup task""" try: if self.cleanup_task is None: self.cleanup_task = asyncio.create_task(self._periodic_cleanup()) logger.debug("Bot provider cleanup task started") except RuntimeError: # No event loop running yet, cleanup will be started later logger.debug("No event loop available for bot provider cleanup task") async def stop_cleanup(self): """Stop the cleanup task""" self._shutdown_event.set() if self.cleanup_task: self.cleanup_task.cancel() try: await self.cleanup_task except asyncio.CancelledError: pass async def _periodic_cleanup(self): """Periodically clean up stale bot providers""" cleanup_interval = 300 # 5 minutes stale_threshold = 900 # 15 minutes while not self._shutdown_event.is_set(): try: await asyncio.sleep(cleanup_interval) now = time.time() providers_to_remove = [] with self.lock: for provider_id, provider in self.bot_providers.items(): if now - provider.last_seen > stale_threshold: providers_to_remove.append(provider_id) logger.info(f"Marking stale bot provider for removal: {provider.name} (ID: {provider_id}, last_seen: {now - provider.last_seen:.1f}s ago)") if providers_to_remove: with self.lock: for provider_id in providers_to_remove: if provider_id in self.bot_providers: del self.bot_providers[provider_id] self._save_bot_providers() logger.info(f"Cleaned up {len(providers_to_remove)} stale bot providers") except asyncio.CancelledError: break except Exception as e: logger.error(f"Error in bot provider cleanup: {e}") def _save_bot_providers(self): """Save bot providers to disk""" try: with self.lock: providers_data = {} for provider_id, provider in self.bot_providers.items(): providers_data[provider_id] = provider.model_dump() with open(self.bot_providers_file, 'w') as f: json.dump(providers_data, f, indent=2) logger.debug(f"Saved {len(providers_data)} bot providers to {self.bot_providers_file}") except Exception as e: logger.error(f"Failed to save bot providers: {e}") def _load_bot_providers(self): """Load bot providers from disk""" try: if not os.path.exists(self.bot_providers_file): logger.debug(f"No bot providers file found at {self.bot_providers_file}") return with open(self.bot_providers_file, 'r') as f: providers_data = json.load(f) with self.lock: for provider_id, provider_dict in providers_data.items(): try: provider = BotProviderModel.model_validate(provider_dict) self.bot_providers[provider_id] = provider except Exception as e: logger.warning(f"Failed to load bot provider {provider_id}: {e}") logger.info(f"Loaded {len(self.bot_providers)} bot providers from {self.bot_providers_file}") except Exception as e: logger.error(f"Failed to load bot providers: {e}") async def register_provider(self, request: BotProviderRegisterRequest) -> BotProviderRegisterResponse: """Register a new bot provider with authentication""" # Check if provider authentication is enabled allowed_providers = BotProviderConfig.get_allowed_providers() if allowed_providers: # Authentication is enabled - validate provider key if request.provider_key not in allowed_providers: logger.warning(f"Rejected bot provider registration with invalid key: {request.provider_key}") raise ValueError("Invalid provider key. Bot provider is not authorized to register.") # Require provider_id for static registration if not request.provider_id: logger.warning("Rejected bot provider registration without provider_id") raise ValueError("provider_id is required for bot provider registration.") provider_id = request.provider_id # Remove any existing provider with the same provider_id (this handles restarts/re-registration) with self.lock: if provider_id in self.bot_providers: logger.info( f"Removing existing provider with ID: {provider_id} for re-registration" ) del self.bot_providers[provider_id] now = time.time() provider = BotProviderModel( provider_id=provider_id, base_url=request.base_url.rstrip("/"), name=request.name, description=request.description, provider_key=request.provider_key, registered_at=now, last_seen=now, ) with self.lock: self.bot_providers[provider_id] = provider # Save to disk self._save_bot_providers() # Start cleanup task if not already running self.start_cleanup() logger.info(f"Registered bot provider: {request.name} at {request.base_url} with key: {request.provider_key}") return BotProviderRegisterResponse(provider_id=provider_id) def list_providers(self) -> BotProviderListResponse: """List all registered bot providers (public information only)""" with self.lock: # Create safe public versions of provider info public_providers: List[BotProviderPublicModel] = [] for provider in self.bot_providers.values(): public_provider = BotProviderPublicModel( provider_id=provider.provider_id, name=provider.name, description=provider.description, registered_at=provider.registered_at, last_seen=provider.last_seen ) public_providers.append(public_provider) return BotProviderListResponse(providers=public_providers) async def list_bots(self) -> BotListResponse: """List all available bots from all registered providers""" bots: List[BotInfoModel] = [] providers: Dict[str, str] = {} # Update last_seen timestamps and fetch bots from each provider with self.lock: providers_copy = dict(self.bot_providers.items()) for provider_id, provider in providers_copy.items(): try: # Update last_seen timestamp with self.lock: if provider_id in self.bot_providers: self.bot_providers[provider_id].last_seen = time.time() # Make HTTP request to provider's /bots endpoint async with httpx.AsyncClient() as client: response = await client.get(f"{provider.base_url}/bots", timeout=5.0) if response.status_code == 200: # Use Pydantic model to validate the response bots_response = BotProviderBotsResponse.model_validate(response.json()) # Add each bot to the consolidated list for bot_info in bots_response.bots: bots.append(bot_info) providers[bot_info.name] = provider_id else: logger.warning(f"Failed to fetch bots from provider {provider.name}: HTTP {response.status_code}") except Exception as e: logger.error(f"Error fetching bots from provider {provider.name}: {e}") continue return BotListResponse(bots=bots, providers=providers) async def get_provider_bots(self, provider_id: str) -> BotProviderBotsResponse: """Get bots from a specific provider""" provider = self.get_provider(provider_id) if not provider: raise ValueError(f"Provider {provider_id} not found") try: async with httpx.AsyncClient() as client: response = await client.get(f"{provider.base_url}/bots", timeout=5.0) if response.status_code == 200: return BotProviderBotsResponse.model_validate(response.json()) else: logger.warning( f"Failed to fetch bots from provider {provider.name}: HTTP {response.status_code}" ) return BotProviderBotsResponse(bots=[]) except Exception as e: logger.error(f"Error fetching bots from provider {provider.name}: {e}") return BotProviderBotsResponse(bots=[]) async def request_bot_join( self, bot_name: str, request: BotJoinLobbyRequest, session_manager: SessionManager, lobby_manager: LobbyManager, ) -> BotJoinLobbyResponse: """Request a bot to join a specific lobby""" # Find which provider has this bot and determine its media capability target_provider_id = request.provider_id bot_has_media = False if not target_provider_id: # Auto-discover provider for this bot with self.lock: providers_copy = dict(self.bot_providers.items()) for provider_id, provider in providers_copy.items(): try: async with httpx.AsyncClient() as client: response = await client.get( f"{provider.base_url}/bots", timeout=5.0 ) if response.status_code == 200: # Use Pydantic model to validate the response bots_response = BotProviderBotsResponse.model_validate( response.json() ) # Look for the bot by name for bot_info in bots_response.bots: if bot_info.name == bot_name: target_provider_id = provider_id bot_has_media = bot_info.has_media break if target_provider_id: break except Exception: continue else: # Query the specified provider for bot media capability with self.lock: if target_provider_id in self.bot_providers: provider = self.bot_providers[target_provider_id] try: async with httpx.AsyncClient() as client: response = await client.get( f"{provider.base_url}/bots", timeout=5.0 ) if response.status_code == 200: # Use Pydantic model to validate the response bots_response = BotProviderBotsResponse.model_validate( response.json() ) # Look for the bot by name for bot_info in bots_response.bots: if bot_info.name == bot_name: bot_has_media = bot_info.has_media break except Exception: # Default to no media if we can't query pass if not target_provider_id: raise ValueError("Bot or provider not found") with self.lock: if target_provider_id not in self.bot_providers: raise ValueError("Provider not found") provider = self.bot_providers[target_provider_id] # Get the lobby to validate it exists lobby = lobby_manager.get_lobby(request.lobby_id) if not lobby: raise ValueError("Lobby not found") # Create a session for the bot bot_session_id = secrets.token_hex(16) bot_instance_id = str(uuid.uuid4()) # Create the Session object for the bot bot_session = session_manager.get_or_create_session( bot_session_id, bot_instance_id=bot_instance_id, has_media=bot_has_media ) logger.info( f"Created bot session for: {bot_session.getName()} (has_media={bot_has_media})" ) # Determine server URL for the bot to connect back to # Use the server's public URL or construct from request server_base_url = os.getenv("PUBLIC_SERVER_URL", "http://localhost:8000") if server_base_url.endswith("/"): server_base_url = server_base_url[:-1] bot_nick = request.nick or f"{bot_name}-bot-{bot_session_id[:8]}" # Get public URL prefix from environment public_url = os.getenv("PUBLIC_URL_PREFIX", "/ai-voicebot") # Prepare the join request for the bot provider bot_join_payload = BotJoinPayload( lobby_id=request.lobby_id, session_id=bot_session_id, nick=bot_nick, server_url=f"{server_base_url}{public_url}".rstrip("/"), insecure=True, # Accept self-signed certificates in development ) try: # Make request to bot provider async with httpx.AsyncClient() as client: response = await client.post( f"{provider.base_url}/bots/{bot_name}/join", json=bot_join_payload.model_dump(), timeout=10.0, ) if response.status_code == 200: # Use Pydantic model to parse and validate response try: join_response = BotProviderJoinResponse.model_validate( response.json() ) run_id = join_response.run_id # Update bot session with run and provider information bot_session.setName(bot_nick) bot_session.bot_run_id = run_id bot_session.bot_provider_id = target_provider_id # Create a unique bot instance ID and track the bot instance bot_instance = BotInstanceModel( bot_instance_id=bot_instance_id, bot_name=bot_name, nick=bot_nick, lobby_id=request.lobby_id, session_id=bot_session_id, provider_id=target_provider_id, run_id=run_id, has_media=bot_has_media, created_at=time.time(), ) # Set the bot_instance_id on the session as well bot_session.bot_instance_id = bot_instance_id with self.lock: self.bot_instances[bot_instance_id] = bot_instance logger.info( f"Bot {bot_name} requested to join lobby {request.lobby_id} with instance ID {bot_instance_id}" ) return BotJoinLobbyResponse( status="requested", bot_instance_id=bot_instance_id, bot_name=bot_name, run_id=run_id, provider_id=target_provider_id, session_id=bot_session_id, ) except ValidationError as e: logger.error(f"Invalid response from bot provider: {e}") raise ValueError( f"Bot provider returned invalid response: {str(e)}" ) else: logger.error( f"Bot provider returned error: HTTP {response.status_code}: {response.text}" ) raise ValueError(f"Bot provider error: {response.status_code}") except httpx.TimeoutException: raise ValueError("Bot provider timeout") except Exception as e: logger.error(f"Error requesting bot join: {e}") raise ValueError(f"Internal server error: {str(e)}") async def request_bot_leave( self, request: BotLeaveLobbyRequest, session_manager: SessionManager ) -> BotLeaveLobbyResponse: """Request a bot to leave from all lobbies and disconnect""" # Find the bot instance with self.lock: if request.bot_instance_id not in self.bot_instances: raise ValueError("Bot instance not found") bot_instance = self.bot_instances[request.bot_instance_id] # Find the bot session bot_session = session_manager.get_session(bot_instance.session_id) if not bot_session: raise ValueError("Bot session not found") if not bot_session.bot_instance_id: raise ValueError("Session is not a bot") logger.info( f"Requesting bot instance {bot_instance.bot_instance_id} to leave all lobbies" ) # Try to stop the bot at the provider level try: with self.lock: if bot_instance.provider_id in self.bot_providers: provider = self.bot_providers[bot_instance.provider_id] try: async with httpx.AsyncClient() as client: response = await client.post( f"{provider.base_url}/bots/runs/{bot_instance.run_id}/stop", timeout=5.0, ) if response.status_code == 200: logger.info( f"Successfully requested bot provider to stop run {bot_instance.run_id}" ) else: logger.warning( f"Bot provider returned error when stopping: HTTP {response.status_code}" ) except Exception as e: logger.warning(f"Failed to request bot stop from provider: {e}") except Exception as e: logger.warning(f"Error communicating with bot provider: {e}") # Force disconnect the bot session from all lobbies try: lobbies_to_leave = bot_session.lobbies[:] for lobby in lobbies_to_leave: try: await bot_session.leave_lobby(lobby) except Exception as e: logger.warning(f"Error removing bot from lobby {lobby.name}: {e}") # Close WebSocket connection if it exists if bot_session.ws: try: await bot_session.ws.close() except Exception as e: logger.warning(f"Error closing bot WebSocket: {e}") bot_session.ws = None except Exception as e: logger.warning(f"Error disconnecting bot session: {e}") # Remove bot instance from tracking with self.lock: if request.bot_instance_id in self.bot_instances: del self.bot_instances[request.bot_instance_id] return BotLeaveLobbyResponse( status="disconnected", bot_instance_id=request.bot_instance_id, session_id=bot_instance.session_id, run_id=bot_instance.run_id, ) async def get_bot_instance(self, bot_instance_id: str) -> BotInstanceModel: """Get information about a specific bot instance""" with self.lock: if bot_instance_id not in self.bot_instances: raise ValueError("Bot instance not found") bot_instance = self.bot_instances[bot_instance_id] return bot_instance def get_bot_instance_id_by_session_id(self, session_id: str) -> Optional[str]: """Get bot_instance_id by session_id""" with self.lock: for bot_instance_id, bot_instance in self.bot_instances.items(): if bot_instance.session_id == session_id: return bot_instance_id return None def get_provider(self, provider_id: str) -> Optional[BotProviderModel]: """Get a specific bot provider by ID""" with self.lock: return self.bot_providers.get(provider_id)