mcpforunityserver 8.2.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (65) hide show
  1. __init__.py +0 -0
  2. core/__init__.py +0 -0
  3. core/config.py +56 -0
  4. core/logging_decorator.py +37 -0
  5. core/telemetry.py +533 -0
  6. core/telemetry_decorator.py +164 -0
  7. main.py +411 -0
  8. mcpforunityserver-8.2.3.dist-info/METADATA +222 -0
  9. mcpforunityserver-8.2.3.dist-info/RECORD +65 -0
  10. mcpforunityserver-8.2.3.dist-info/WHEEL +5 -0
  11. mcpforunityserver-8.2.3.dist-info/entry_points.txt +2 -0
  12. mcpforunityserver-8.2.3.dist-info/licenses/LICENSE +21 -0
  13. mcpforunityserver-8.2.3.dist-info/top_level.txt +8 -0
  14. models/__init__.py +4 -0
  15. models/models.py +56 -0
  16. models/unity_response.py +47 -0
  17. routes/__init__.py +0 -0
  18. services/__init__.py +0 -0
  19. services/custom_tool_service.py +339 -0
  20. services/registry/__init__.py +22 -0
  21. services/registry/resource_registry.py +53 -0
  22. services/registry/tool_registry.py +51 -0
  23. services/resources/__init__.py +81 -0
  24. services/resources/active_tool.py +47 -0
  25. services/resources/custom_tools.py +57 -0
  26. services/resources/editor_state.py +42 -0
  27. services/resources/layers.py +29 -0
  28. services/resources/menu_items.py +34 -0
  29. services/resources/prefab_stage.py +39 -0
  30. services/resources/project_info.py +39 -0
  31. services/resources/selection.py +55 -0
  32. services/resources/tags.py +30 -0
  33. services/resources/tests.py +55 -0
  34. services/resources/unity_instances.py +122 -0
  35. services/resources/windows.py +47 -0
  36. services/tools/__init__.py +76 -0
  37. services/tools/batch_execute.py +78 -0
  38. services/tools/debug_request_context.py +71 -0
  39. services/tools/execute_custom_tool.py +38 -0
  40. services/tools/execute_menu_item.py +29 -0
  41. services/tools/find_in_file.py +174 -0
  42. services/tools/manage_asset.py +129 -0
  43. services/tools/manage_editor.py +63 -0
  44. services/tools/manage_gameobject.py +240 -0
  45. services/tools/manage_material.py +95 -0
  46. services/tools/manage_prefabs.py +62 -0
  47. services/tools/manage_scene.py +75 -0
  48. services/tools/manage_script.py +602 -0
  49. services/tools/manage_shader.py +64 -0
  50. services/tools/read_console.py +115 -0
  51. services/tools/run_tests.py +108 -0
  52. services/tools/script_apply_edits.py +998 -0
  53. services/tools/set_active_instance.py +112 -0
  54. services/tools/utils.py +60 -0
  55. transport/__init__.py +0 -0
  56. transport/legacy/port_discovery.py +329 -0
  57. transport/legacy/stdio_port_registry.py +65 -0
  58. transport/legacy/unity_connection.py +785 -0
  59. transport/models.py +62 -0
  60. transport/plugin_hub.py +412 -0
  61. transport/plugin_registry.py +123 -0
  62. transport/unity_instance_middleware.py +141 -0
  63. transport/unity_transport.py +103 -0
  64. utils/module_discovery.py +55 -0
  65. utils/reload_sentinel.py +9 -0
@@ -0,0 +1,785 @@
1
+ from core.config import config
2
+ import contextlib
3
+ from dataclasses import dataclass
4
+ import errno
5
+ import json
6
+ import logging
7
+ import os
8
+ from pathlib import Path
9
+ from transport.legacy.port_discovery import PortDiscovery
10
+ import random
11
+ import socket
12
+ import struct
13
+ import threading
14
+ import time
15
+ from typing import Any
16
+
17
+ from models.models import MCPResponse, UnityInstanceInfo
18
+ from transport.legacy.stdio_port_registry import stdio_port_registry
19
+
20
+
21
+ logger = logging.getLogger("mcp-for-unity-server")
22
+
23
+ # Module-level lock to guard global connection initialization
24
+ _connection_lock = threading.Lock()
25
+
26
+ # Maximum allowed framed payload size (64 MiB)
27
+ FRAMED_MAX = 64 * 1024 * 1024
28
+
29
+
30
+ @dataclass
31
+ class UnityConnection:
32
+ """Manages the socket connection to the Unity Editor."""
33
+ host: str = config.unity_host
34
+ port: int = None # Will be set dynamically
35
+ sock: socket.socket = None # Socket for Unity communication
36
+ use_framing: bool = False # Negotiated per-connection
37
+ instance_id: str | None = None # Instance identifier for reconnection
38
+
39
+ def __post_init__(self):
40
+ """Set port from discovery if not explicitly provided"""
41
+ if self.port is None:
42
+ self.port = stdio_port_registry.get_port(self.instance_id)
43
+ self._io_lock = threading.Lock()
44
+ self._conn_lock = threading.Lock()
45
+
46
+ def _prepare_socket(self, sock: socket.socket) -> None:
47
+ try:
48
+ sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
49
+ except OSError as exc:
50
+ logger.debug(f"Unable to set TCP_NODELAY: {exc}")
51
+
52
+ def connect(self) -> bool:
53
+ """Establish a connection to the Unity Editor."""
54
+ if self.sock:
55
+ return True
56
+ with self._conn_lock:
57
+ if self.sock:
58
+ return True
59
+ try:
60
+ # Bounded connect to avoid indefinite blocking
61
+ connect_timeout = float(
62
+ getattr(config, "connection_timeout", 1.0))
63
+ # We trust config.unity_host (default 127.0.0.1) but future improvements
64
+ # could dynamically prefer 'localhost' depending on OS resolver behavior.
65
+ self.sock = socket.create_connection(
66
+ (self.host, self.port), connect_timeout)
67
+ self._prepare_socket(self.sock)
68
+ logger.debug(f"Connected to Unity at {self.host}:{self.port}")
69
+
70
+ # Strict handshake: require FRAMING=1
71
+ try:
72
+ require_framing = getattr(config, "require_framing", True)
73
+ handshake_timeout = float(
74
+ getattr(config, "handshake_timeout", 1.0))
75
+ self.sock.settimeout(handshake_timeout)
76
+ buf = bytearray()
77
+ deadline = time.monotonic() + handshake_timeout
78
+ while time.monotonic() < deadline and len(buf) < 512:
79
+ try:
80
+ chunk = self.sock.recv(256)
81
+ if not chunk:
82
+ break
83
+ buf.extend(chunk)
84
+ if b"\n" in buf:
85
+ break
86
+ except socket.timeout:
87
+ break
88
+ text = bytes(buf).decode('ascii', errors='ignore').strip()
89
+
90
+ if 'FRAMING=1' in text:
91
+ self.use_framing = True
92
+ logger.debug(
93
+ 'MCP for Unity handshake received: FRAMING=1 (strict)')
94
+ else:
95
+ if require_framing:
96
+ # Best-effort plain-text advisory for legacy peers
97
+ with contextlib.suppress(Exception):
98
+ self.sock.sendall(
99
+ b'MCP for Unity requires FRAMING=1\n')
100
+ raise ConnectionError(
101
+ f'MCP for Unity requires FRAMING=1, got: {text!r}')
102
+ else:
103
+ self.use_framing = False
104
+ logger.warning(
105
+ 'MCP for Unity handshake missing FRAMING=1; proceeding in legacy mode by configuration')
106
+ finally:
107
+ self.sock.settimeout(config.connection_timeout)
108
+ return True
109
+ except Exception as e:
110
+ logger.error(f"Failed to connect to Unity: {str(e)}")
111
+ try:
112
+ if self.sock:
113
+ self.sock.close()
114
+ except Exception:
115
+ pass
116
+ self.sock = None
117
+ return False
118
+
119
+ def disconnect(self):
120
+ """Close the connection to the Unity Editor."""
121
+ if self.sock:
122
+ try:
123
+ self.sock.close()
124
+ except Exception as e:
125
+ logger.error(f"Error disconnecting from Unity: {str(e)}")
126
+ finally:
127
+ self.sock = None
128
+
129
+ def _read_exact(self, sock: socket.socket, count: int) -> bytes:
130
+ data = bytearray()
131
+ while len(data) < count:
132
+ chunk = sock.recv(count - len(data))
133
+ if not chunk:
134
+ raise ConnectionError(
135
+ "Connection closed before reading expected bytes")
136
+ data.extend(chunk)
137
+ return bytes(data)
138
+
139
+ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes:
140
+ """Receive a complete response from Unity, handling chunked data."""
141
+ if self.use_framing:
142
+ # Heartbeat semantics: the Unity editor emits zero-length frames while
143
+ # a long-running command is still executing. We tolerate a bounded
144
+ # number of these frames (or a small time window) before surfacing a
145
+ # timeout to the caller so tools can retry or fail gracefully.
146
+ heartbeat_limit = getattr(config, 'max_heartbeat_frames', 16)
147
+ heartbeat_window = getattr(config, 'heartbeat_timeout', 2.0)
148
+ heartbeat_started = time.monotonic()
149
+ heartbeat_count = 0
150
+ try:
151
+ while True:
152
+ header = self._read_exact(sock, 8)
153
+ payload_len = struct.unpack('>Q', header)[0]
154
+ if payload_len == 0:
155
+ heartbeat_count += 1
156
+ logger.debug(
157
+ f"Received heartbeat frame #{heartbeat_count}")
158
+ if heartbeat_count >= heartbeat_limit or (time.monotonic() - heartbeat_started) > heartbeat_window:
159
+ raise TimeoutError(
160
+ "Unity sent heartbeat frames without payload within configured threshold"
161
+ )
162
+ continue
163
+ if payload_len > FRAMED_MAX:
164
+ raise ValueError(
165
+ f"Invalid framed length: {payload_len}")
166
+ payload = self._read_exact(sock, payload_len)
167
+ logger.debug(
168
+ f"Received framed response ({len(payload)} bytes)")
169
+ return payload
170
+ except socket.timeout as exc:
171
+ logger.warning("Socket timeout during framed receive")
172
+ raise TimeoutError("Timeout receiving Unity response") from exc
173
+ except TimeoutError:
174
+ raise
175
+ except Exception as exc:
176
+ logger.error(f"Error during framed receive: {exc}")
177
+ raise
178
+
179
+ chunks = []
180
+ # Respect the socket's currently configured timeout
181
+ try:
182
+ while True:
183
+ chunk = sock.recv(buffer_size)
184
+ if not chunk:
185
+ if not chunks:
186
+ raise Exception(
187
+ "Connection closed before receiving data")
188
+ break
189
+ chunks.append(chunk)
190
+
191
+ # Process the data received so far
192
+ data = b''.join(chunks)
193
+ decoded_data = data.decode('utf-8')
194
+
195
+ # Check if we've received a complete response
196
+ try:
197
+ # Special case for ping-pong
198
+ if decoded_data.strip().startswith('{"status":"success","result":{"message":"pong"'):
199
+ logger.debug("Received ping response")
200
+ return data
201
+
202
+ # Handle escaped quotes in the content
203
+ if '"content":' in decoded_data:
204
+ # Find the content field and its value
205
+ content_start = decoded_data.find('"content":') + 9
206
+ content_end = decoded_data.rfind('"', content_start)
207
+ if content_end > content_start:
208
+ # Replace escaped quotes in content with regular quotes
209
+ content = decoded_data[content_start:content_end]
210
+ content = content.replace('\\"', '"')
211
+ decoded_data = decoded_data[:content_start] + \
212
+ content + decoded_data[content_end:]
213
+
214
+ # Validate JSON format
215
+ json.loads(decoded_data)
216
+
217
+ # If we get here, we have valid JSON
218
+ logger.info(
219
+ f"Received complete response ({len(data)} bytes)")
220
+ return data
221
+ except json.JSONDecodeError:
222
+ # We haven't received a complete valid JSON response yet
223
+ continue
224
+ except Exception as e:
225
+ logger.warning(
226
+ f"Error processing response chunk: {str(e)}")
227
+ # Continue reading more chunks as this might not be the complete response
228
+ continue
229
+ except socket.timeout:
230
+ logger.warning("Socket timeout during receive")
231
+ raise Exception("Timeout receiving Unity response")
232
+ except Exception as e:
233
+ logger.error(f"Error during receive: {str(e)}")
234
+ raise
235
+
236
+ def send_command(self, command_type: str, params: dict[str, Any] = None) -> dict[str, Any]:
237
+ """Send a command with retry/backoff and port rediscovery. Pings only when requested."""
238
+ # Defensive guard: catch empty/placeholder invocations early
239
+ if not command_type:
240
+ raise ValueError("MCP call missing command_type")
241
+ if params is None:
242
+ return MCPResponse(success=False, error="MCP call received with no parameters (client placeholder?)")
243
+ attempts = max(config.max_retries, 5)
244
+ base_backoff = max(0.5, config.retry_delay)
245
+
246
+ def read_status_file(target_hash: str | None = None) -> dict | None:
247
+ try:
248
+ base_path = Path.home().joinpath('.unity-mcp')
249
+ status_files = sorted(
250
+ base_path.glob('unity-mcp-status-*.json'),
251
+ key=lambda p: p.stat().st_mtime,
252
+ reverse=True,
253
+ )
254
+ if not status_files:
255
+ return None
256
+ if target_hash:
257
+ for status_path in status_files:
258
+ if status_path.stem.endswith(target_hash):
259
+ with status_path.open('r') as f:
260
+ return json.load(f)
261
+ # Fallback: return most recent regardless of hash
262
+ with status_files[0].open('r') as f:
263
+ return json.load(f)
264
+ except FileNotFoundError:
265
+ logger.debug(
266
+ "Unity status file disappeared before it could be read")
267
+ return None
268
+ except json.JSONDecodeError as exc:
269
+ logger.warning(f"Malformed Unity status file: {exc}")
270
+ return None
271
+ except OSError as exc:
272
+ logger.warning(f"Failed to read Unity status file: {exc}")
273
+ return None
274
+ except Exception as exc:
275
+ logger.debug(f"Preflight status check failed: {exc}")
276
+ return None
277
+
278
+ last_short_timeout = None
279
+
280
+ # Extract hash suffix from instance id (e.g., Project@hash)
281
+ target_hash: str | None = None
282
+ if self.instance_id and '@' in self.instance_id:
283
+ maybe_hash = self.instance_id.split('@', 1)[1].strip()
284
+ if maybe_hash:
285
+ target_hash = maybe_hash
286
+
287
+ # Preflight: if Unity reports reloading, return a structured hint so clients can retry politely
288
+ try:
289
+ status = read_status_file(target_hash)
290
+ if status and (status.get('reloading') or status.get('reason') == 'reloading'):
291
+ return MCPResponse(
292
+ success=False,
293
+ error="Unity is reloading; please retry",
294
+ hint="retry",
295
+ )
296
+ except Exception as exc:
297
+ logger.debug(f"Preflight status check failed: {exc}")
298
+
299
+ for attempt in range(attempts + 1):
300
+ try:
301
+ # Ensure connected (handshake occurs within connect())
302
+ if not self.sock and not self.connect():
303
+ raise ConnectionError("Could not connect to Unity")
304
+
305
+ # Build payload
306
+ if command_type == 'ping':
307
+ payload = b'ping'
308
+ else:
309
+ payload = json.dumps({
310
+ 'type': command_type,
311
+ 'params': params,
312
+ }).encode('utf-8')
313
+
314
+ # Send/receive are serialized to protect the shared socket
315
+ with self._io_lock:
316
+ mode = 'framed' if self.use_framing else 'legacy'
317
+ with contextlib.suppress(Exception):
318
+ logger.debug(
319
+ f"send {len(payload)} bytes; mode={mode}; head={payload[:32].decode('utf-8', 'ignore')}")
320
+ if self.use_framing:
321
+ header = struct.pack('>Q', len(payload))
322
+ self.sock.sendall(header)
323
+ self.sock.sendall(payload)
324
+ else:
325
+ self.sock.sendall(payload)
326
+
327
+ # During retry bursts use a short receive timeout and ensure restoration
328
+ restore_timeout = None
329
+ if attempt > 0 and last_short_timeout is None:
330
+ restore_timeout = self.sock.gettimeout()
331
+ self.sock.settimeout(1.0)
332
+ try:
333
+ response_data = self.receive_full_response(self.sock)
334
+ with contextlib.suppress(Exception):
335
+ logger.debug(
336
+ f"recv {len(response_data)} bytes; mode={mode}")
337
+ finally:
338
+ if restore_timeout is not None:
339
+ self.sock.settimeout(restore_timeout)
340
+ last_short_timeout = None
341
+
342
+ # Parse
343
+ if command_type == 'ping':
344
+ resp = json.loads(response_data.decode('utf-8'))
345
+ if resp.get('status') == 'success' and resp.get('result', {}).get('message') == 'pong':
346
+ return {"message": "pong"}
347
+ raise Exception("Ping unsuccessful")
348
+
349
+ resp = json.loads(response_data.decode('utf-8'))
350
+ if resp.get('status') == 'error':
351
+ err = resp.get('error') or resp.get(
352
+ 'message', 'Unknown Unity error')
353
+ raise Exception(err)
354
+ return resp.get('result', {})
355
+ except Exception as e:
356
+ logger.warning(
357
+ f"Unity communication attempt {attempt+1} failed: {e}")
358
+ try:
359
+ if self.sock:
360
+ self.sock.close()
361
+ finally:
362
+ self.sock = None
363
+
364
+ # Re-discover the port for this specific instance
365
+ try:
366
+ new_port: int | None = None
367
+ if self.instance_id:
368
+ # Try to rediscover the specific instance via shared registry
369
+ refreshed_instance = stdio_port_registry.get_instance(
370
+ self.instance_id)
371
+ if refreshed_instance and isinstance(refreshed_instance.port, int):
372
+ new_port = refreshed_instance.port
373
+ logger.debug(
374
+ f"Rediscovered instance {self.instance_id} on port {new_port}")
375
+ else:
376
+ logger.warning(
377
+ f"Instance {self.instance_id} not found during reconnection; falling back to port scan",
378
+ )
379
+
380
+ # Fallback to registry default if instance-specific discovery failed
381
+ if new_port is None:
382
+ new_port = stdio_port_registry.get_port(
383
+ self.instance_id)
384
+ logger.info(
385
+ f"Using Unity port from stdio_port_registry: {new_port}")
386
+
387
+ if new_port != self.port:
388
+ logger.info(
389
+ f"Unity port changed {self.port} -> {new_port}")
390
+ self.port = new_port
391
+ except Exception as de:
392
+ logger.debug(f"Port discovery failed: {de}")
393
+
394
+ if attempt < attempts:
395
+ # Heartbeat-aware, jittered backoff
396
+ status = read_status_file(target_hash)
397
+ # Base exponential backoff
398
+ backoff = base_backoff * (2 ** attempt)
399
+ # Decorrelated jitter multiplier
400
+ jitter = random.uniform(0.1, 0.3)
401
+
402
+ # Fast‑retry for transient socket failures
403
+ fast_error = isinstance(
404
+ e, (ConnectionRefusedError, ConnectionResetError, TimeoutError))
405
+ if not fast_error:
406
+ try:
407
+ err_no = getattr(e, 'errno', None)
408
+ fast_error = err_no in (
409
+ errno.ECONNREFUSED, errno.ECONNRESET, errno.ETIMEDOUT)
410
+ except Exception:
411
+ pass
412
+
413
+ # Cap backoff depending on state
414
+ if status and status.get('reloading'):
415
+ cap = 0.8
416
+ elif fast_error:
417
+ cap = 0.25
418
+ else:
419
+ cap = 3.0
420
+
421
+ sleep_s = min(cap, jitter * (2 ** attempt))
422
+ time.sleep(sleep_s)
423
+ continue
424
+ raise
425
+
426
+
427
+ # -----------------------------
428
+ # Connection Pool for Multiple Unity Instances
429
+ # -----------------------------
430
+
431
+ class UnityConnectionPool:
432
+ """Manages connections to multiple Unity Editor instances"""
433
+
434
+ def __init__(self):
435
+ self._connections: dict[str, UnityConnection] = {}
436
+ self._known_instances: dict[str, UnityInstanceInfo] = {}
437
+ self._last_full_scan: float = 0
438
+ self._scan_interval: float = 5.0 # Cache for 5 seconds
439
+ self._pool_lock = threading.Lock()
440
+ self._default_instance_id: str | None = None
441
+
442
+ # Check for default instance from environment
443
+ env_default = os.environ.get("UNITY_MCP_DEFAULT_INSTANCE", "").strip()
444
+ if env_default:
445
+ self._default_instance_id = env_default
446
+ logger.info(
447
+ f"Default Unity instance set from environment: {env_default}")
448
+
449
+ def discover_all_instances(self, force_refresh: bool = False) -> list[UnityInstanceInfo]:
450
+ """
451
+ Discover all running Unity Editor instances.
452
+
453
+ Args:
454
+ force_refresh: If True, bypass cache and scan immediately
455
+
456
+ Returns:
457
+ List of UnityInstanceInfo objects
458
+ """
459
+ now = time.time()
460
+
461
+ # Return cached results if valid
462
+ if not force_refresh and (now - self._last_full_scan) < self._scan_interval:
463
+ logger.debug(
464
+ f"Returning cached Unity instances (age: {now - self._last_full_scan:.1f}s)")
465
+ return list(self._known_instances.values())
466
+
467
+ # Scan for instances
468
+ logger.debug("Scanning for Unity instances...")
469
+ instances = PortDiscovery.discover_all_unity_instances()
470
+
471
+ # Update cache
472
+ with self._pool_lock:
473
+ self._known_instances = {inst.id: inst for inst in instances}
474
+ self._last_full_scan = now
475
+
476
+ logger.info(
477
+ f"Found {len(instances)} Unity instances: {[inst.id for inst in instances]}")
478
+ return instances
479
+
480
+ def _resolve_instance_id(self, instance_identifier: str | None, instances: list[UnityInstanceInfo]) -> UnityInstanceInfo:
481
+ """
482
+ Resolve an instance identifier to a specific Unity instance.
483
+
484
+ Args:
485
+ instance_identifier: User-provided identifier (name, hash, name@hash, path, port, or None)
486
+ instances: List of available instances
487
+
488
+ Returns:
489
+ Resolved UnityInstanceInfo
490
+
491
+ Raises:
492
+ ConnectionError: If instance cannot be resolved
493
+ """
494
+ if not instances:
495
+ raise ConnectionError(
496
+ "No Unity Editor instances found. Please ensure Unity is running with MCP for Unity bridge."
497
+ )
498
+
499
+ # Use default instance if no identifier provided
500
+ if instance_identifier is None:
501
+ if self._default_instance_id:
502
+ instance_identifier = self._default_instance_id
503
+ logger.debug(f"Using default instance: {instance_identifier}")
504
+ else:
505
+ # Use the most recently active instance
506
+ # Instances with no heartbeat (None) should be sorted last (use 0 as sentinel)
507
+ sorted_instances = sorted(
508
+ instances,
509
+ key=lambda inst: inst.last_heartbeat.timestamp() if inst.last_heartbeat else 0.0,
510
+ reverse=True,
511
+ )
512
+ logger.info(
513
+ f"No instance specified, using most recent: {sorted_instances[0].id}")
514
+ return sorted_instances[0]
515
+
516
+ identifier = instance_identifier.strip()
517
+
518
+ # Try exact ID match first
519
+ for inst in instances:
520
+ if inst.id == identifier:
521
+ return inst
522
+
523
+ # Try project name match
524
+ name_matches = [inst for inst in instances if inst.name == identifier]
525
+ if len(name_matches) == 1:
526
+ return name_matches[0]
527
+ elif len(name_matches) > 1:
528
+ # Multiple projects with same name - return helpful error
529
+ suggestions = [
530
+ {
531
+ "id": inst.id,
532
+ "path": inst.path,
533
+ "port": inst.port,
534
+ "suggest": f"Use unity_instance='{inst.id}'"
535
+ }
536
+ for inst in name_matches
537
+ ]
538
+ raise ConnectionError(
539
+ f"Project name '{identifier}' matches {len(name_matches)} instances. "
540
+ f"Please use the full format (e.g., '{name_matches[0].id}'). "
541
+ f"Available instances: {suggestions}"
542
+ )
543
+
544
+ # Try hash match
545
+ hash_matches = [inst for inst in instances if inst.hash ==
546
+ identifier or inst.hash.startswith(identifier)]
547
+ if len(hash_matches) == 1:
548
+ return hash_matches[0]
549
+ elif len(hash_matches) > 1:
550
+ raise ConnectionError(
551
+ f"Hash '{identifier}' matches multiple instances: {[inst.id for inst in hash_matches]}"
552
+ )
553
+
554
+ # Try composite format: Name@Hash or Name@Port
555
+ if "@" in identifier:
556
+ name_part, hint_part = identifier.split("@", 1)
557
+ composite_matches = [
558
+ inst for inst in instances
559
+ if inst.name == name_part and (
560
+ inst.hash.startswith(hint_part) or str(
561
+ inst.port) == hint_part
562
+ )
563
+ ]
564
+ if len(composite_matches) == 1:
565
+ return composite_matches[0]
566
+
567
+ # Try port match (as string)
568
+ try:
569
+ port_num = int(identifier)
570
+ port_matches = [
571
+ inst for inst in instances if inst.port == port_num]
572
+ if len(port_matches) == 1:
573
+ return port_matches[0]
574
+ except ValueError:
575
+ pass
576
+
577
+ # Try path match
578
+ path_matches = [inst for inst in instances if inst.path == identifier]
579
+ if len(path_matches) == 1:
580
+ return path_matches[0]
581
+
582
+ # Nothing matched
583
+ available_ids = [inst.id for inst in instances]
584
+ raise ConnectionError(
585
+ f"Unity instance '{identifier}' not found. "
586
+ f"Available instances: {available_ids}. "
587
+ f"Check unity://instances resource for all instances."
588
+ )
589
+
590
+ def get_connection(self, instance_identifier: str | None = None) -> UnityConnection:
591
+ """
592
+ Get or create a connection to a Unity instance.
593
+
594
+ Args:
595
+ instance_identifier: Optional identifier (name, hash, name@hash, etc.)
596
+ If None, uses default or most recent instance
597
+
598
+ Returns:
599
+ UnityConnection to the specified instance
600
+
601
+ Raises:
602
+ ConnectionError: If instance cannot be found or connected
603
+ """
604
+ # Refresh instance list if cache expired
605
+ instances = self.discover_all_instances()
606
+
607
+ # Resolve identifier to specific instance
608
+ target = self._resolve_instance_id(instance_identifier, instances)
609
+
610
+ # Return existing connection or create new one
611
+ with self._pool_lock:
612
+ if target.id not in self._connections:
613
+ logger.info(
614
+ f"Creating new connection to Unity instance: {target.id} (port {target.port})")
615
+ conn = UnityConnection(port=target.port, instance_id=target.id)
616
+ if not conn.connect():
617
+ raise ConnectionError(
618
+ f"Failed to connect to Unity instance '{target.id}' on port {target.port}. "
619
+ f"Ensure the Unity Editor is running."
620
+ )
621
+ self._connections[target.id] = conn
622
+ else:
623
+ # Update existing connection with instance_id and port if changed
624
+ conn = self._connections[target.id]
625
+ conn.instance_id = target.id
626
+ if conn.port != target.port:
627
+ logger.info(
628
+ f"Updating cached port for {target.id}: {conn.port} -> {target.port}")
629
+ conn.port = target.port
630
+ logger.debug(f"Reusing existing connection to: {target.id}")
631
+
632
+ return self._connections[target.id]
633
+
634
+ def disconnect_all(self):
635
+ """Disconnect all active connections"""
636
+ with self._pool_lock:
637
+ for instance_id, conn in self._connections.items():
638
+ try:
639
+ logger.info(
640
+ f"Disconnecting from Unity instance: {instance_id}")
641
+ conn.disconnect()
642
+ except Exception:
643
+ logger.exception(f"Error disconnecting from {instance_id}")
644
+ self._connections.clear()
645
+
646
+
647
+ # Global Unity connection pool
648
+ _unity_connection_pool: UnityConnectionPool | None = None
649
+ _pool_init_lock = threading.Lock()
650
+
651
+
652
+ def get_unity_connection_pool() -> UnityConnectionPool:
653
+ """Get or create the global Unity connection pool"""
654
+ global _unity_connection_pool
655
+
656
+ if _unity_connection_pool is not None:
657
+ return _unity_connection_pool
658
+
659
+ with _pool_init_lock:
660
+ if _unity_connection_pool is not None:
661
+ return _unity_connection_pool
662
+
663
+ logger.info("Initializing Unity connection pool")
664
+ _unity_connection_pool = UnityConnectionPool()
665
+ return _unity_connection_pool
666
+
667
+
668
+ # Backwards compatibility: keep old single-connection function
669
+ def get_unity_connection(instance_identifier: str | None = None) -> UnityConnection:
670
+ """Retrieve or establish a Unity connection.
671
+
672
+ Args:
673
+ instance_identifier: Optional identifier for specific Unity instance.
674
+ If None, uses default or most recent instance.
675
+
676
+ Returns:
677
+ UnityConnection to the specified or default Unity instance
678
+
679
+ Note: This function now uses the connection pool internally.
680
+ """
681
+ pool = get_unity_connection_pool()
682
+ return pool.get_connection(instance_identifier)
683
+
684
+
685
+ # -----------------------------
686
+ # Centralized retry helpers
687
+ # -----------------------------
688
+
689
+ def _is_reloading_response(resp: object) -> bool:
690
+ """Return True if the Unity response indicates the editor is reloading.
691
+
692
+ Supports both raw dict payloads from Unity and MCPResponse objects returned
693
+ by preflight checks or transport helpers.
694
+ """
695
+ # Structured MCPResponse from preflight/transport
696
+ if isinstance(resp, MCPResponse):
697
+ # Explicit "please retry" hint from preflight
698
+ if getattr(resp, "hint", None) == "retry":
699
+ return True
700
+ message_text = f"{resp.message or ''} {resp.error or ''}".lower()
701
+ return "reload" in message_text
702
+
703
+ # Raw Unity payloads
704
+ if isinstance(resp, dict):
705
+ if resp.get("state") == "reloading":
706
+ return True
707
+ message_text = (resp.get("message") or resp.get("error") or "").lower()
708
+ return "reload" in message_text
709
+
710
+ return False
711
+
712
+
713
+ def send_command_with_retry(
714
+ command_type: str,
715
+ params: dict[str, Any],
716
+ *,
717
+ instance_id: str | None = None,
718
+ max_retries: int | None = None,
719
+ retry_ms: int | None = None
720
+ ) -> dict[str, Any] | MCPResponse:
721
+ """Send a command to a Unity instance, waiting politely through Unity reloads.
722
+
723
+ Args:
724
+ command_type: The command type to send
725
+ params: Command parameters
726
+ instance_id: Optional Unity instance identifier (name, hash, name@hash, etc.)
727
+ max_retries: Maximum number of retries for reload states
728
+ retry_ms: Delay between retries in milliseconds
729
+
730
+ Returns:
731
+ Response dictionary or MCPResponse from Unity
732
+
733
+ Uses config.reload_retry_ms and config.reload_max_retries by default. Preserves the
734
+ structured failure if retries are exhausted.
735
+ """
736
+ conn = get_unity_connection(instance_id)
737
+ if max_retries is None:
738
+ max_retries = getattr(config, "reload_max_retries", 40)
739
+ if retry_ms is None:
740
+ retry_ms = getattr(config, "reload_retry_ms", 250)
741
+
742
+ response = conn.send_command(command_type, params)
743
+ retries = 0
744
+ while _is_reloading_response(response) and retries < max_retries:
745
+ delay_ms = int(response.get("retry_after_ms", retry_ms)
746
+ ) if isinstance(response, dict) else retry_ms
747
+ time.sleep(max(0.0, delay_ms / 1000.0))
748
+ retries += 1
749
+ response = conn.send_command(command_type, params)
750
+ return response
751
+
752
+
753
+ async def async_send_command_with_retry(
754
+ command_type: str,
755
+ params: dict[str, Any],
756
+ *,
757
+ instance_id: str | None = None,
758
+ loop=None,
759
+ max_retries: int | None = None,
760
+ retry_ms: int | None = None
761
+ ) -> dict[str, Any] | MCPResponse:
762
+ """Async wrapper that runs the blocking retry helper in a thread pool.
763
+
764
+ Args:
765
+ command_type: The command type to send
766
+ params: Command parameters
767
+ instance_id: Optional Unity instance identifier
768
+ loop: Optional asyncio event loop
769
+ max_retries: Maximum number of retries for reload states
770
+ retry_ms: Delay between retries in milliseconds
771
+
772
+ Returns:
773
+ Response dictionary or MCPResponse on error
774
+ """
775
+ try:
776
+ import asyncio # local import to avoid mandatory asyncio dependency for sync callers
777
+ if loop is None:
778
+ loop = asyncio.get_running_loop()
779
+ return await loop.run_in_executor(
780
+ None,
781
+ lambda: send_command_with_retry(
782
+ command_type, params, instance_id=instance_id, max_retries=max_retries, retry_ms=retry_ms),
783
+ )
784
+ except Exception as e:
785
+ return MCPResponse(success=False, error=str(e))