|
| 1 | +""" |
| 2 | +Connection pool for managing multiple Unity Editor instances. |
| 3 | +""" |
| 4 | +import logging |
| 5 | +import os |
| 6 | +import threading |
| 7 | +import time |
| 8 | + |
| 9 | +from models import UnityInstanceInfo |
| 10 | +from port_discovery import PortDiscovery |
| 11 | + |
| 12 | +logger = logging.getLogger(__name__) |
| 13 | + |
| 14 | + |
| 15 | +class UnityConnectionPool: |
| 16 | + """Manages connections to multiple Unity Editor instances""" |
| 17 | + |
| 18 | + def __init__(self): |
| 19 | + # Import here to avoid circular dependency |
| 20 | + from unity_connection import UnityConnection |
| 21 | + self._UnityConnection = UnityConnection |
| 22 | + |
| 23 | + self._connections: dict[str, "UnityConnection"] = {} |
| 24 | + self._known_instances: dict[str, UnityInstanceInfo] = {} |
| 25 | + self._last_full_scan: float = 0 |
| 26 | + self._scan_interval: float = 5.0 # Cache for 5 seconds |
| 27 | + self._pool_lock = threading.Lock() |
| 28 | + self._default_instance_id: str | None = None |
| 29 | + |
| 30 | + # Check for default instance from environment |
| 31 | + env_default = os.environ.get("UNITY_MCP_DEFAULT_INSTANCE", "").strip() |
| 32 | + if env_default: |
| 33 | + self._default_instance_id = env_default |
| 34 | + logger.info(f"Default Unity instance set from environment: {env_default}") |
| 35 | + |
| 36 | + def discover_all_instances(self, force_refresh: bool = False) -> list[UnityInstanceInfo]: |
| 37 | + """ |
| 38 | + Discover all running Unity Editor instances. |
| 39 | +
|
| 40 | + Args: |
| 41 | + force_refresh: If True, bypass cache and scan immediately |
| 42 | +
|
| 43 | + Returns: |
| 44 | + List of UnityInstanceInfo objects |
| 45 | + """ |
| 46 | + now = time.time() |
| 47 | + |
| 48 | + # Return cached results if valid |
| 49 | + if not force_refresh and (now - self._last_full_scan) < self._scan_interval: |
| 50 | + logger.debug(f"Returning cached Unity instances (age: {now - self._last_full_scan:.1f}s)") |
| 51 | + return list(self._known_instances.values()) |
| 52 | + |
| 53 | + # Scan for instances |
| 54 | + logger.debug("Scanning for Unity instances...") |
| 55 | + instances = PortDiscovery.discover_all_unity_instances() |
| 56 | + |
| 57 | + # Update cache |
| 58 | + with self._pool_lock: |
| 59 | + self._known_instances = {inst.id: inst for inst in instances} |
| 60 | + self._last_full_scan = now |
| 61 | + |
| 62 | + logger.info(f"Found {len(instances)} Unity instances: {[inst.id for inst in instances]}") |
| 63 | + return instances |
| 64 | + |
| 65 | + def _resolve_instance_id(self, instance_identifier: str | None, instances: list[UnityInstanceInfo]) -> UnityInstanceInfo: |
| 66 | + """ |
| 67 | + Resolve an instance identifier to a specific Unity instance. |
| 68 | +
|
| 69 | + Args: |
| 70 | + instance_identifier: User-provided identifier (name, hash, name@hash, path, port, or None) |
| 71 | + instances: List of available instances |
| 72 | +
|
| 73 | + Returns: |
| 74 | + Resolved UnityInstanceInfo |
| 75 | +
|
| 76 | + Raises: |
| 77 | + ConnectionError: If instance cannot be resolved |
| 78 | + """ |
| 79 | + if not instances: |
| 80 | + raise ConnectionError( |
| 81 | + "No Unity Editor instances found. Please ensure Unity is running with MCP for Unity bridge." |
| 82 | + ) |
| 83 | + |
| 84 | + # Use default instance if no identifier provided |
| 85 | + if instance_identifier is None: |
| 86 | + if self._default_instance_id: |
| 87 | + instance_identifier = self._default_instance_id |
| 88 | + logger.debug(f"Using default instance: {instance_identifier}") |
| 89 | + else: |
| 90 | + # Use the most recently active instance |
| 91 | + # Instances with no heartbeat (None) should be sorted last (use 0.0 as sentinel) |
| 92 | + sorted_instances = sorted( |
| 93 | + instances, |
| 94 | + key=lambda inst: inst.last_heartbeat.timestamp() if inst.last_heartbeat else 0.0, |
| 95 | + reverse=True, |
| 96 | + ) |
| 97 | + logger.info(f"No instance specified, using most recent: {sorted_instances[0].id}") |
| 98 | + return sorted_instances[0] |
| 99 | + |
| 100 | + identifier = instance_identifier.strip() |
| 101 | + |
| 102 | + # Try exact ID match first |
| 103 | + for inst in instances: |
| 104 | + if inst.id == identifier: |
| 105 | + return inst |
| 106 | + |
| 107 | + # Try project name match |
| 108 | + name_matches = [inst for inst in instances if inst.name == identifier] |
| 109 | + if len(name_matches) == 1: |
| 110 | + return name_matches[0] |
| 111 | + elif len(name_matches) > 1: |
| 112 | + # Multiple projects with same name - return helpful error |
| 113 | + suggestions = [ |
| 114 | + { |
| 115 | + "id": inst.id, |
| 116 | + "path": inst.path, |
| 117 | + "port": inst.port, |
| 118 | + "suggest": f"Use unity_instance='{inst.id}'" |
| 119 | + } |
| 120 | + for inst in name_matches |
| 121 | + ] |
| 122 | + raise ConnectionError( |
| 123 | + f"Project name '{identifier}' matches {len(name_matches)} instances. " |
| 124 | + f"Please use the full format (e.g., '{name_matches[0].id}'). " |
| 125 | + f"Available instances: {suggestions}" |
| 126 | + ) |
| 127 | + |
| 128 | + # Try hash match |
| 129 | + hash_matches = [inst for inst in instances if inst.hash == identifier or inst.hash.startswith(identifier)] |
| 130 | + if len(hash_matches) == 1: |
| 131 | + return hash_matches[0] |
| 132 | + elif len(hash_matches) > 1: |
| 133 | + raise ConnectionError( |
| 134 | + f"Hash '{identifier}' matches multiple instances: {[inst.id for inst in hash_matches]}" |
| 135 | + ) |
| 136 | + |
| 137 | + # Try composite format: Name@Hash or Name@Port |
| 138 | + if "@" in identifier: |
| 139 | + name_part, hint_part = identifier.split("@", 1) |
| 140 | + composite_matches = [ |
| 141 | + inst for inst in instances |
| 142 | + if inst.name == name_part and ( |
| 143 | + inst.hash.startswith(hint_part) or str(inst.port) == hint_part |
| 144 | + ) |
| 145 | + ] |
| 146 | + if len(composite_matches) == 1: |
| 147 | + return composite_matches[0] |
| 148 | + |
| 149 | + # Try port match (as string) |
| 150 | + try: |
| 151 | + port_num = int(identifier) |
| 152 | + port_matches = [inst for inst in instances if inst.port == port_num] |
| 153 | + if len(port_matches) == 1: |
| 154 | + return port_matches[0] |
| 155 | + except ValueError: |
| 156 | + pass |
| 157 | + |
| 158 | + # Try path match |
| 159 | + path_matches = [inst for inst in instances if inst.path == identifier] |
| 160 | + if len(path_matches) == 1: |
| 161 | + return path_matches[0] |
| 162 | + |
| 163 | + # Nothing matched |
| 164 | + available_ids = [inst.id for inst in instances] |
| 165 | + raise ConnectionError( |
| 166 | + f"Unity instance '{identifier}' not found. " |
| 167 | + f"Available instances: {available_ids}. " |
| 168 | + f"Use the unity_instances resource to see all instances." |
| 169 | + ) |
| 170 | + |
| 171 | + def get_connection(self, instance_identifier: str | None = None): |
| 172 | + """ |
| 173 | + Get or create a connection to a Unity instance. |
| 174 | +
|
| 175 | + Args: |
| 176 | + instance_identifier: Optional identifier (name, hash, name@hash, etc.) |
| 177 | + If None, uses default or most recent instance |
| 178 | +
|
| 179 | + Returns: |
| 180 | + UnityConnection to the specified instance |
| 181 | +
|
| 182 | + Raises: |
| 183 | + ConnectionError: If instance cannot be found or connected |
| 184 | + """ |
| 185 | + # Refresh instance list if cache expired |
| 186 | + instances = self.discover_all_instances() |
| 187 | + |
| 188 | + # Resolve identifier to specific instance |
| 189 | + target = self._resolve_instance_id(instance_identifier, instances) |
| 190 | + |
| 191 | + # Return existing connection or create new one |
| 192 | + with self._pool_lock: |
| 193 | + if target.id not in self._connections: |
| 194 | + logger.info(f"Creating new connection to Unity instance: {target.id} (port {target.port})") |
| 195 | + conn = self._UnityConnection(port=target.port, instance_id=target.id) |
| 196 | + if not conn.connect(): |
| 197 | + raise ConnectionError( |
| 198 | + f"Failed to connect to Unity instance '{target.id}' on port {target.port}. " |
| 199 | + f"Ensure the Unity Editor is running." |
| 200 | + ) |
| 201 | + self._connections[target.id] = conn |
| 202 | + else: |
| 203 | + # Update existing connection with instance_id and port if changed |
| 204 | + conn = self._connections[target.id] |
| 205 | + conn.instance_id = target.id |
| 206 | + if conn.port != target.port: |
| 207 | + logger.info(f"Updating cached port for {target.id}: {conn.port} -> {target.port}") |
| 208 | + conn.port = target.port |
| 209 | + logger.debug(f"Reusing existing connection to: {target.id}") |
| 210 | + |
| 211 | + return self._connections[target.id] |
| 212 | + |
| 213 | + def disconnect_all(self): |
| 214 | + """Disconnect all active connections""" |
| 215 | + with self._pool_lock: |
| 216 | + for instance_id, conn in self._connections.items(): |
| 217 | + try: |
| 218 | + logger.info(f"Disconnecting from Unity instance: {instance_id}") |
| 219 | + conn.disconnect() |
| 220 | + except Exception: |
| 221 | + logger.exception(f"Error disconnecting from {instance_id}") |
| 222 | + self._connections.clear() |
| 223 | + |
| 224 | + |
| 225 | +# Global Unity connection pool |
| 226 | +_unity_connection_pool: UnityConnectionPool | None = None |
| 227 | +_pool_init_lock = threading.Lock() |
| 228 | + |
| 229 | + |
| 230 | +def get_unity_connection_pool() -> UnityConnectionPool: |
| 231 | + """Get or create the global Unity connection pool.""" |
| 232 | + global _unity_connection_pool |
| 233 | + if _unity_connection_pool is None: |
| 234 | + with _pool_init_lock: |
| 235 | + if _unity_connection_pool is None: |
| 236 | + _unity_connection_pool = UnityConnectionPool() |
| 237 | + return _unity_connection_pool |
0 commit comments