diff --git a/Server/src/core/config.py b/Server/src/core/config.py index 4639f9e68..c7c6460f8 100644 --- a/Server/src/core/config.py +++ b/Server/src/core/config.py @@ -32,6 +32,8 @@ class ServerConfig: # Connection settings connection_timeout: float = 30.0 + # Hard ceiling on a command's total time across all retries (wedged-socket guard). + command_total_timeout: float = 90.0 buffer_size: int = 16 * 1024 * 1024 # 16MB buffer # STDIO framing behaviour diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index 1d703a862..42b9d498e 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -50,7 +50,7 @@ def _prepare_socket(self, sock: socket.socket) -> None: except OSError as exc: logger.debug(f"Unable to set TCP_NODELAY: {exc}") - def connect(self) -> bool: + def connect(self, connect_timeout: float | None = None) -> bool: """Establish a connection to the Unity Editor.""" if self.sock: return True @@ -59,8 +59,9 @@ def connect(self) -> bool: return True try: # Bounded connect to avoid indefinite blocking - connect_timeout = float( - getattr(config, "connection_timeout", 1.0)) + if connect_timeout is None: + connect_timeout = float( + getattr(config, "connection_timeout", 1.0)) # We trust config.unity_host (default 127.0.0.1) but future improvements # could dynamically prefer 'localhost' depending on OS resolver behavior. self.sock = socket.create_connection( @@ -129,12 +130,13 @@ def disconnect(self): self.sock = None def _ensure_live_connection(self) -> None: - """Detect and discard stale (peer-closed) sockets before sending. + """Detect and discard cleanly-closed sockets before sending. - After domain reload Unity closes all TCP connections. The Python side - may still hold a reference to the dead socket. A non-blocking peek - detects this so send_command can reconnect instead of writing to a dead - socket and getting 'Connection closed before reading expected bytes'. + A graceful domain reload closes the bridge's TCP connections (FIN), so a + non-blocking peek sees EOF and lets send_command reconnect instead of + writing to a dead socket. A half-open socket left without a FIN is not + caught here; that case is bounded by the per-command deadline and the + reloading-status preflight. """ if not self.sock: return @@ -265,13 +267,20 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: logger.error(f"Error during receive: {str(e)}") raise - def send_command(self, command_type: str, params: dict[str, Any] = None, max_attempts: int | None = None) -> dict[str, Any]: + def _cap_to_deadline(self, timeout: float, deadline: float | None, floor: float = 0.05) -> float: + """Shrink a blocking timeout to whatever budget remains before the deadline.""" + if deadline is None: + return timeout + return max(floor, min(timeout, deadline - time.monotonic())) + + def send_command(self, command_type: str, params: dict[str, Any] | None = None, max_attempts: int | None = None, deadline: float | None = None) -> dict[str, Any]: """Send a command with retry/backoff and port rediscovery. Pings only when requested. Args: command_type: The Unity command to send params: Command parameters max_attempts: Maximum retry attempts (None = use config default, 0 = no retries) + deadline: Shared monotonic() ceiling across retries (None = derive from command_total_timeout) """ # Defensive guard: catch empty/placeholder invocations early if not command_type: @@ -282,6 +291,11 @@ def send_command(self, command_type: str, params: dict[str, Any] = None, max_att 5) if max_attempts is None else max_attempts base_backoff = max(0.5, config.retry_delay) + # Cap total time across all retries so a wedged socket can't block unbounded. + total_timeout = max(0.0, float(getattr(config, "command_total_timeout", 90.0))) + if deadline is None and total_timeout > 0: + deadline = time.monotonic() + total_timeout + def read_status_file(target_hash: str | None = None) -> dict | None: try: base_path = Path.home().joinpath('.unity-mcp') @@ -314,8 +328,6 @@ def read_status_file(target_hash: str | None = None) -> dict | None: logger.debug(f"Preflight status check failed: {exc}") return None - last_short_timeout = None - # Extract hash suffix from instance id (e.g., Project@hash) target_hash: str | None = None if self.instance_id and '@' in self.instance_id: @@ -327,6 +339,10 @@ def read_status_file(target_hash: str | None = None) -> dict | None: try: status = read_status_file(target_hash) if status and (status.get('reloading') or status.get('reason') == 'reloading'): + # Reload invalidates the socket; drop it under the I/O lock so this + # close is serialized against the send/recv block, then reconnect next call. + with self._io_lock: + self.disconnect() return MCPResponse( success=False, error="Unity is reloading; please retry", @@ -336,13 +352,20 @@ def read_status_file(target_hash: str | None = None) -> dict | None: logger.debug(f"Preflight status check failed: {exc}") for attempt in range(attempts + 1): + if deadline is not None and time.monotonic() >= deadline: + logger.warning( + "Command '%s' exceeded total deadline of %.1fs after %d attempt(s); giving up", + command_type, total_timeout, attempt) + raise TimeoutError( + f"Command '{command_type}' exceeded total deadline of " + f"{total_timeout:.1f}s (connection wedged or Unity unresponsive)") try: # Discard stale sockets left over from a previous domain reload # so we reconnect instead of writing to a dead connection. self._ensure_live_connection() # Ensure connected (handshake occurs within connect()) t_conn_start = time.time() - if not self.sock and not self.connect(): + if not self.sock and not self.connect(self._cap_to_deadline(config.connection_timeout, deadline)): raise ConnectionError("Could not connect to Unity") logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type) @@ -370,11 +393,17 @@ def read_status_file(target_hash: str | None = None) -> dict | None: self.sock.sendall(payload) logger.info("[TIMING-STDIO] sendall took %.3fs command=%s", time.time() - t_send_start, command_type) - # During retry bursts use a short receive timeout and ensure restoration + # Cap the receive timeout to the remaining command budget (and use a + # short timeout during retry bursts) so a wedged socket can't block + # past the deadline. restore_timeout = None - if attempt > 0 and last_short_timeout is None: + recv_timeout = 1.0 if attempt > 0 else self.sock.gettimeout() + if deadline is not None: + recv_timeout = self._cap_to_deadline( + recv_timeout or config.connection_timeout, deadline) + if recv_timeout is not None and recv_timeout != self.sock.gettimeout(): restore_timeout = self.sock.gettimeout() - self.sock.settimeout(1.0) + self.sock.settimeout(recv_timeout) try: t_recv_start = time.time() response_data = self.receive_full_response(self.sock) @@ -385,7 +414,6 @@ def read_status_file(target_hash: str | None = None) -> dict | None: finally: if restore_timeout is not None: self.sock.settimeout(restore_timeout) - last_short_timeout = None # Parse if command_type == 'ping': @@ -468,6 +496,7 @@ def read_status_file(target_hash: str | None = None) -> dict | None: cap = 3.0 sleep_s = min(cap, jitter * (2 ** attempt)) + sleep_s = self._cap_to_deadline(sleep_s, deadline, floor=0.0) time.sleep(sleep_s) continue raise @@ -834,16 +863,19 @@ def send_command_with_retry( # Clamp to [0, 20] to prevent misconfiguration from causing excessive waits max_wait_s = max(0.0, min(max_wait_s, 20.0)) + total_timeout = max(0.0, float(getattr(config, "command_total_timeout", 90.0))) + deadline = time.monotonic() + total_timeout if total_timeout > 0 else None + # If retry_on_reload=False, disable connection-level retries too (issue #577) # Commands that trigger compilation/reload shouldn't retry on disconnect send_max_attempts = None if retry_on_reload else 0 response = conn.send_command( - command_type, params, max_attempts=send_max_attempts) + command_type, params, max_attempts=send_max_attempts, deadline=deadline) retries = 0 wait_started = None reason = _extract_response_reason(response) - while retry_on_reload and _is_reloading_response(response) and retries < max_retries: + while retry_on_reload and _is_reloading_response(response) and retries < max_retries and (deadline is None or time.monotonic() < deadline): if wait_started is None: wait_started = time.monotonic() logger.debug( @@ -876,7 +908,7 @@ def send_command_with_retry( ) time.sleep(max(0.0, sleep_ms / 1000.0)) retries += 1 - response = conn.send_command(command_type, params) + response = conn.send_command(command_type, params, deadline=deadline) reason = _extract_response_reason(response) if wait_started is not None: diff --git a/Server/tests/integration/test_connection_deadline.py b/Server/tests/integration/test_connection_deadline.py new file mode 100644 index 000000000..d51e6bf55 --- /dev/null +++ b/Server/tests/integration/test_connection_deadline.py @@ -0,0 +1,110 @@ +"""Connection resilience when a domain reload leaves the Unity socket half-open.""" +from __future__ import annotations + +import json +import socket +import threading +import time +from pathlib import Path + +import pytest + +from core.config import config +import transport.legacy.unity_connection as uc +from transport.legacy.unity_connection import UnityConnection + + +def _start_silent_bridge(): + """Accept and handshake, then never answer (a wedged half-open bridge).""" + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(8) + port = srv.getsockname()[1] + accepted: list[socket.socket] = [] + stop = threading.Event() + + def serve(): + srv.settimeout(0.5) + while not stop.is_set(): + try: + client, _ = srv.accept() + except socket.timeout: + continue + except OSError: + break + accepted.append(client) + try: + client.sendall(b"WELCOME UNITY-MCP 1 FRAMING=1\n") + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + return srv, port, stop, accepted + + +def _write_reloading_status(tmp_path): + status_dir = Path(tmp_path) / ".unity-mcp" + status_dir.mkdir(parents=True, exist_ok=True) + (status_dir / "unity-mcp-status-deadbeef.json").write_text( + json.dumps({"reloading": True, "reason": "reloading", "project_path": "x"}) + ) + + +@pytest.fixture +def silent_bridge(monkeypatch, tmp_path): + srv, port, stop, accepted = _start_silent_bridge() + monkeypatch.setattr(uc.Path, "home", lambda: Path(tmp_path)) + monkeypatch.setattr(uc.stdio_port_registry, "get_port", lambda instance_id=None: port) + monkeypatch.setattr(uc.stdio_port_registry, "get_instance", lambda instance_id: None) + monkeypatch.setattr(config, "connection_timeout", 1.0) + monkeypatch.setattr(config, "command_total_timeout", 2.0, raising=False) + try: + yield port + finally: + stop.set() + try: + srv.close() + except OSError: + pass + for client in accepted: + try: + client.close() + except OSError: + pass + + +def test_send_command_against_wedged_socket_is_bounded(silent_bridge, monkeypatch): + """A connection_timeout longer than the total budget must not let a single blocking + recv overrun the ceiling: the deadline caps each recv, so the call still stops near + command_total_timeout rather than connection_timeout.""" + monkeypatch.setattr(config, "connection_timeout", 5.0) + conn = UnityConnection(port=silent_bridge, instance_id="Repro@deadbeef") + start = time.monotonic() + with pytest.raises(TimeoutError, match="exceeded total deadline"): + conn.send_command("get_editor_state", {}) + assert time.monotonic() - start < 4.0 + + +def test_send_command_with_retry_is_bounded_by_deadline(silent_bridge, tmp_path, monkeypatch): + """The reload-wait loop honours command_total_timeout, not just max_wait_s.""" + conn = UnityConnection(port=silent_bridge, instance_id="Repro@deadbeef") + monkeypatch.setattr(uc, "get_unity_connection", lambda instance_id=None: conn) + _write_reloading_status(tmp_path) + + start = time.monotonic() + resp = uc.send_command_with_retry("get_editor_state", {}, instance_id="Repro@deadbeef") + assert time.monotonic() - start < 6.0 + assert uc._extract_response_reason(resp) == "reloading" + + +def test_reload_signal_drops_socket_for_fresh_reconnect(silent_bridge, tmp_path): + """A 'reloading' status drops the socket so the next command reconnects.""" + conn = UnityConnection(port=silent_bridge, instance_id="Repro@deadbeef") + assert conn.connect() is True + assert conn.sock is not None + _write_reloading_status(tmp_path) + + response = conn.send_command("get_editor_state", {}) + assert conn.sock is None + assert uc._extract_response_reason(response) == "reloading"