From ef1be2c3dcf65f40ce463728150b6173677afc12 Mon Sep 17 00:00:00 2001 From: Marco Meyer Date: Tue, 11 Aug 2026 10:03:55 +0900 Subject: [PATCH 1/9] fix(transport): make Unity timeouts long-budget, loud, and unmistakable Long-running custom tools (captures, image comparison) exceeded the hub's default 30s command budget and callers got the five words {success:false, error:"TimeoutError", hint:null} - after which agents carried on reading stale on-disk artifacts as if the tool had run. - custom tools now default to timeout_seconds=300 (explicit caller values win; hub honors up to 3600) - CustomToolService._normalize_response no longer drops the hint field - transport-level timeouts now say what a timeout means: the command may or may not have executed, verify effects once editor state reports ready_for_tools, and pass timeout_seconds for long tools; structured data carries reason=unity_timeout - _poll_until_complete timeout message gains the same warning + retry hint --- Server/src/services/custom_tool_service.py | 14 +++- Server/src/transport/unity_transport.py | 17 +++- Server/tests/test_custom_tool_timeouts.py | 96 ++++++++++++++++++++++ 3 files changed, 125 insertions(+), 2 deletions(-) create mode 100644 Server/tests/test_custom_tool_timeouts.py diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index fe24bc7725..5414c84e8b 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -152,6 +152,12 @@ async def execute_tool( message=f"Tool '{tool_name}' not found for project {project_id}", ) + # Custom tools are frequently long-running (captures, image comparison, + # batch imports), so the hub's default 30s command budget times them out + # while they are still working. Give them a generous default; callers can + # still override with an explicit timeout_seconds parameter (up to 3600). + params.setdefault("timeout_seconds", 300) + response = await send_with_unity_instance( async_send_command_with_retry, unity_instance, @@ -216,8 +222,13 @@ async def _poll_until_complete( if time.time() > deadline: return MCPResponse( success=False, - message=f"Timeout waiting for {tool_name} to complete", + message=( + f"Timeout waiting for {tool_name} to complete. The operation may " + "still be running in Unity — do not assume it finished or failed; " + "verify its effects before continuing." + ), data=self._safe_response(response), + hint="retry", ) await asyncio.sleep(poll_interval) @@ -281,6 +292,7 @@ def _normalize_response(self, response) -> MCPResponse: error=response.get("error"), data=response.get( "data", response) if "data" not in response else response["data"], + hint=response.get("hint"), ) success = True diff --git a/Server/src/transport/unity_transport.py b/Server/src/transport/unity_transport.py index 9d07f4714c..57149c3d9c 100644 --- a/Server/src/transport/unity_transport.py +++ b/Server/src/transport/unity_transport.py @@ -1,6 +1,7 @@ """Transport helpers for routing commands to Unity.""" from __future__ import annotations +import asyncio import logging from typing import Awaitable, Callable, TypeVar @@ -101,11 +102,25 @@ async def send_with_unity_instance( except Exception as exc: # NOTE: asyncio.TimeoutError has an empty str() by default, which is confusing for clients. err = str(exc) or f"{type(exc).__name__}" + data: dict | None = None + if isinstance(exc, (asyncio.TimeoutError, TimeoutError)): + # A bare "TimeoutError" reads as noise and agents plough on as if the + # command ran. Spell out what a timeout means and what to do about it. + err = ( + f"Unity did not respond to '{command_type}' within its time budget. " + "The editor is likely compiling, reloading, importing assets, or still " + "running the operation. The command may or may not have executed — do NOT " + "assume it ran; verify its effects (or re-run it) once " + "mcpforunity://editor/state reports advice.ready_for_tools=true. " + "For long-running tools, pass timeout_seconds in the tool parameters " + "(honored up to 3600)." + ) + data = {"reason": "unity_timeout"} # Fail fast with a retry hint instead of hanging for COMMAND_TIMEOUT. # The client can decide whether retrying is appropriate for the command. return normalize_unity_response( MCPResponse(success=False, error=err, - hint="retry").model_dump() + hint="retry", data=data).model_dump() ) if unity_instance: diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py new file mode 100644 index 0000000000..5e970febc8 --- /dev/null +++ b/Server/tests/test_custom_tool_timeouts.py @@ -0,0 +1,96 @@ +"""Timeouts on custom tools must be long-budget, loud, and impossible to misread. + +Real-world failure (2026-08-11): a long-running project tool (image comparison) +exceeded the hub's default 30s command budget and the caller got the five words +{"success":false,"error":"TimeoutError","hint":null} — after which the agent +carried on reading stale on-disk artifacts as if the tool had run. Three fixes: +custom tools default to a 300s budget, _normalize_response no longer drops the +hint, and transport-level timeouts explain that the command may not have run. +""" + +import asyncio + +import pytest + +from core.config import config +from models.models import MCPResponse, ToolDefinitionModel +from services.custom_tool_service import CustomToolService + + +class _NoopMcp: + def custom_route(self, _path, methods=None): # noqa: ARG002 + def _decorator(fn): + return fn + + return _decorator + + def tool(self, name=None, description=None): # noqa: ARG002 + def _decorator(fn): + return fn + + return _decorator + + +def _make_service() -> CustomToolService: + return CustomToolService(_NoopMcp(), project_scoped_tools=True) + + +def test_normalize_response_preserves_hint(): + service = _make_service() + result = service._normalize_response( + {"success": False, "error": "boom", "hint": "retry", "data": {"reason": "x"}} + ) + assert isinstance(result, MCPResponse) + assert result.hint == "retry" + assert result.data == {"reason": "x"} + + +@pytest.mark.asyncio +async def test_execute_tool_defaults_long_timeout(monkeypatch): + service = _make_service() + definition = ToolDefinitionModel(name="long_running_tool", description="long tool") + service._register_tool("proj-1", definition) + + seen_params: dict = {} + + async def fake_send(_send_fn, _instance, _command, params, **kwargs): # noqa: ARG001 + seen_params.update(params) + return {"success": True, "message": "ok"} + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", fake_send) + + result = await service.execute_tool("proj-1", "long_running_tool", "Playground@abc", {}) + assert result.success is True + assert seen_params["timeout_seconds"] == 300 + + # An explicit caller value must win over the default. + seen_params.clear() + await service.execute_tool( + "proj-1", "long_running_tool", "Playground@abc", {"timeout_seconds": 900}) + assert seen_params["timeout_seconds"] == 900 + + +@pytest.mark.asyncio +async def test_transport_timeout_error_is_self_explanatory(monkeypatch): + from transport import unity_transport + + monkeypatch.setattr(config, "transport_mode", "http") + + async def raise_timeout(*args, **kwargs): # noqa: ARG001 + raise asyncio.TimeoutError() + + monkeypatch.setattr( + unity_transport.PluginHub, "send_command_for_instance", raise_timeout) + + async def unused_send_fn(*args, **kwargs): # noqa: ARG001 + raise AssertionError("stdio path should not be used in http mode") + + response = await unity_transport.send_with_unity_instance( + unused_send_fn, "Playground@abc", "long_running_tool", {}) + + assert response["success"] is False + assert "may or may not have executed" in response["error"] + assert "timeout_seconds" in response["error"] + assert response["hint"] == "retry" + assert response["data"] == {"reason": "unity_timeout"} From 813da1958b47945ecc57b0d8bff0a912ccb05880 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Tue, 27 Oct 2026 15:09:51 +0900 Subject: [PATCH 2/9] fix(custom-tools): keep an explicit timeoutSeconds; no retry hint at the polling deadline Review follow-up. - The 300 s default was added with setdefault("timeout_seconds"), and the hub reads timeout_seconds before timeoutSeconds, so a caller's explicit timeoutSeconds lost to the default. The default is now added only when neither spelling is present. - The polling-deadline response no longer carries hint="retry": the operation may still be running, and resubmitting it could run it twice. The message still says to verify its effects first. Tests: an explicit timeoutSeconds is passed through untouched; the polling-deadline response says the operation may still be running and has no hint. --- Server/src/services/custom_tool_service.py | 8 ++++--- Server/tests/test_custom_tool_timeouts.py | 26 ++++++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index 5414c84e8b..e71fce4988 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -155,8 +155,10 @@ async def execute_tool( # Custom tools are frequently long-running (captures, image comparison, # batch imports), so the hub's default 30s command budget times them out # while they are still working. Give them a generous default; callers can - # still override with an explicit timeout_seconds parameter (up to 3600). - params.setdefault("timeout_seconds", 300) + # still override with an explicit timeout_seconds (or timeoutSeconds) parameter, up to 3600. + # The hub reads timeout_seconds first, so a default added beside a caller's timeoutSeconds would win. + if "timeout_seconds" not in params and "timeoutSeconds" not in params: + params["timeout_seconds"] = 300 response = await send_with_unity_instance( async_send_command_with_retry, @@ -228,7 +230,7 @@ async def _poll_until_complete( "verify its effects before continuing." ), data=self._safe_response(response), - hint="retry", + # No retry hint: the operation may still be running, and resubmitting it could run it twice. ) await asyncio.sleep(poll_interval) diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index 5e970febc8..3193da0ac4 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -70,6 +70,32 @@ async def fake_send(_send_fn, _instance, _command, params, **kwargs): # noqa: A "proj-1", "long_running_tool", "Playground@abc", {"timeout_seconds": 900}) assert seen_params["timeout_seconds"] == 900 + # So must the camelCase spelling the hub also honors: no snake_case default beside it. + seen_params.clear() + await service.execute_tool( + "proj-1", "long_running_tool", "Playground@abc", {"timeoutSeconds": 1200}) + assert seen_params == {"timeoutSeconds": 1200} + + +@pytest.mark.asyncio +async def test_polling_deadline_does_not_suggest_resubmitting(monkeypatch): + """At the polling deadline the operation may still run: say so, and do not hint a retry.""" + service = _make_service() + + async def still_running(_send_fn, _instance, _command, params, **kwargs): # noqa: ARG001 + return {"_mcp_status": "pending", "message": "working"} + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", still_running) + + result = await service._poll_until_complete( + "long_running_tool", "Playground@abc", {}, {"_mcp_status": "pending"}, "status", + max_poll_seconds=0.01) + + assert result.success is False + assert "may still be running" in result.message + assert result.hint is None + @pytest.mark.asyncio async def test_transport_timeout_error_is_self_explanatory(monkeypatch): From 79b20c978981e160eabe8c9ac146ae4ee6faaf59 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Tue, 27 Oct 2026 15:14:13 +0900 Subject: [PATCH 3/9] fix(custom-tools): give each status poll only what is left of the polling budget Review follow-up. A poll copies the tool's parameters, including the 300 s default command budget, so one status poll could wait far past max_poll_seconds. Each poll now asks for the remaining budget (at least 1 s; any timeoutSeconds copy is dropped). The hub's 30 s minimum wait is unchanged: it applies to every command, and it bounds a poll's overrun to 30 s. Test: a poll's timeout_seconds is within the remaining budget. --- Server/src/services/custom_tool_service.py | 6 ++++++ Server/tests/test_custom_tool_timeouts.py | 23 ++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index e71fce4988..78b773ba03 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -235,6 +235,12 @@ async def _poll_until_complete( await asyncio.sleep(poll_interval) + # A poll inherits the tool's parameters, including its long command budget (300 s by default): + # left as is, one status poll could wait far past this polling deadline. Give each poll what is + # left of the budget instead (the hub still waits at least its 30 s default). + poll_params.pop("timeoutSeconds", None) + poll_params["timeout_seconds"] = max(1.0, deadline - time.time()) + try: response = await send_with_unity_instance( async_send_command_with_retry, diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index 3193da0ac4..bac2a26b28 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -120,3 +120,26 @@ async def unused_send_fn(*args, **kwargs): # noqa: ARG001 assert "timeout_seconds" in response["error"] assert response["hint"] == "retry" assert response["data"] == {"reason": "unity_timeout"} + + +@pytest.mark.asyncio +async def test_each_poll_waits_no_longer_than_the_remaining_budget(monkeypatch): + """A poll must not inherit the tool's 300 s command budget past the polling deadline.""" + service = _make_service() + seen: list[dict] = [] + + async def complete(_send_fn, _instance, _command, params, **kwargs): # noqa: ARG001 + seen.append(dict(params)) + return {"_mcp_status": "complete", "success": True, "message": "done"} + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", complete) + + result = await service._poll_until_complete( + "long_running_tool", "Playground@abc", {"timeout_seconds": 300, "timeoutSeconds": 300}, + {"_mcp_status": "pending", "_mcp_poll_interval": 0.01}, "status", max_poll_seconds=5) + + assert result.success is True + assert seen and 0 < seen[0]["timeout_seconds"] <= 5 + assert "timeoutSeconds" not in seen[0] + assert seen[0]["action"] == "status" From ce12e6b1d6f613096dfb55ebccf5d6fe0bd077c1 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Thu, 29 Oct 2026 22:28:04 +0900 Subject: [PATCH 4/9] fix(custom-tools): cut a status poll off at the polling deadline Each poll already asks for only the remaining budget, but the hub waits at least its 30 s default for any command, so a poll with 1 s left could still wait 30 s. The poll is now bounded with asyncio.wait_for on the remaining budget (the hub drops the pending command when cancelled), and the sleep before it is clamped to the deadline too. Test: a hung poll with max_poll_seconds=1 ends in under 2 s (30.1 s before). --- Server/src/services/custom_tool_service.py | 27 ++++++++++++++-------- Server/tests/test_custom_tool_timeouts.py | 23 ++++++++++++++++++ 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index 78b773ba03..0e4579b46f 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -233,22 +233,31 @@ async def _poll_until_complete( # No retry hint: the operation may still be running, and resubmitting it could run it twice. ) - await asyncio.sleep(poll_interval) + await asyncio.sleep(min(poll_interval, max(0.0, deadline - time.time()))) # A poll inherits the tool's parameters, including its long command budget (300 s by default): # left as is, one status poll could wait far past this polling deadline. Give each poll what is - # left of the budget instead (the hub still waits at least its 30 s default). + # left of the budget instead. The hub still waits at least its 30 s default for any command, + # so the wait itself is also cut off at the deadline here. + remaining = max(0.0, deadline - time.time()) poll_params.pop("timeoutSeconds", None) - poll_params["timeout_seconds"] = max(1.0, deadline - time.time()) + poll_params["timeout_seconds"] = max(1.0, remaining) try: - response = await send_with_unity_instance( - async_send_command_with_retry, - unity_instance, - tool_name, - poll_params, - user_id=user_id, + response = await asyncio.wait_for( + send_with_unity_instance( + async_send_command_with_retry, + unity_instance, + tool_name, + poll_params, + user_id=user_id, + ), + timeout=remaining, ) + except asyncio.TimeoutError: + # The deadline passed mid-poll: keep the last pending response and let the + # check above report the timeout. + continue except Exception as exc: # pragma: no cover - network/domain reload variability logger.debug(f"Polling {tool_name} failed, will retry: {exc}") # Back off modestly but stay responsive. diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index bac2a26b28..808b6ff153 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -143,3 +143,26 @@ async def complete(_send_fn, _instance, _command, params, **kwargs): # noqa: AR assert seen and 0 < seen[0]["timeout_seconds"] <= 5 assert "timeoutSeconds" not in seen[0] assert seen[0]["action"] == "status" + + +@pytest.mark.asyncio +async def test_a_hung_poll_is_cut_off_at_the_polling_deadline(monkeypatch): + """The hub waits at least 30 s for any command; a poll must still end at the polling deadline.""" + service = _make_service() + + async def hang(_send_fn, _instance, _command, _params, **kwargs): # noqa: ARG001 + await asyncio.sleep(30) + + monkeypatch.setattr( + "services.custom_tool_service.send_with_unity_instance", hang) + + loop = asyncio.get_running_loop() + started = loop.time() + result = await service._poll_until_complete( + "long_running_tool", "Playground@abc", {}, + {"_mcp_status": "pending", "_mcp_poll_interval": 0.01}, "status", max_poll_seconds=1) + elapsed = loop.time() - started + + assert result.success is False + assert "Timeout waiting for long_running_tool" in result.message + assert elapsed < 2.0 From f59888fa056ff25295f95f06e43efaeaca3d55af Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Thu, 29 Oct 2026 22:50:00 +0900 Subject: [PATCH 5/9] fix(custom-tools): a stdio status poll ends by itself at the polling deadline On stdio a poll runs send_command_with_retry in an executor thread. asyncio.wait_for stops waiting at the deadline but cannot stop that thread, which kept going until command_total_timeout (90 s by default) while holding the connection's I/O lock. send_command_with_retry and async_send_command_with_retry take an optional budget_s that can only shorten that deadline (the socket connect/receive timeouts and the reload wait are already capped by it), and each poll passes what is left of its budget. A poll is skipped once nothing is left. Tests: with a blocking stdio connection and max_poll_seconds=1, the poll thread gets a deadline within the budget and ends by it (before: a 90 s deadline, still running when the poll returned); budget_s shortens but never extends the configured deadline. --- Server/src/services/custom_tool_service.py | 7 ++- .../src/transport/legacy/unity_connection.py | 14 ++++- Server/tests/test_custom_tool_timeouts.py | 60 +++++++++++++++++++ 3 files changed, 77 insertions(+), 4 deletions(-) diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index 0e4579b46f..8c15152f93 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -238,8 +238,12 @@ async def _poll_until_complete( # A poll inherits the tool's parameters, including its long command budget (300 s by default): # left as is, one status poll could wait far past this polling deadline. Give each poll what is # left of the budget instead. The hub still waits at least its 30 s default for any command, - # so the wait itself is also cut off at the deadline here. + # so the wait itself is also cut off at the deadline here. On stdio the poll runs in a worker + # thread that a cancelled wait cannot stop, so it is given the same budget (budget_s) and + # ends by itself, releasing the connection. remaining = max(0.0, deadline - time.time()) + if remaining <= 0: + continue poll_params.pop("timeoutSeconds", None) poll_params["timeout_seconds"] = max(1.0, remaining) @@ -251,6 +255,7 @@ async def _poll_until_complete( tool_name, poll_params, user_id=user_id, + budget_s=remaining, ), timeout=remaining, ) diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index 986781bed0..1440fa50d4 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -818,7 +818,8 @@ def send_command_with_retry( instance_id: str | None = None, max_retries: int | None = None, retry_ms: int | None = None, - retry_on_reload: bool = True + retry_on_reload: bool = True, + budget_s: float | None = None, ) -> dict[str, Any] | MCPResponse: """Send a command to a Unity instance, waiting politely through Unity reloads. @@ -830,6 +831,8 @@ def send_command_with_retry( retry_ms: Delay between retries in milliseconds retry_on_reload: If False, don't retry when Unity is reloading (for commands that trigger compilation/reload and shouldn't be re-sent) + budget_s: Optional cap in seconds on the whole call, connection retries and reload + waits included; it can only shorten config.command_total_timeout Returns: Response dictionary or MCPResponse from Unity @@ -870,6 +873,9 @@ def send_command_with_retry( 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 budget_s is not None: + capped = time.monotonic() + max(0.0, budget_s) + deadline = capped if deadline is None else min(deadline, capped) # If retry_on_reload=False, disable connection-level retries too (issue #577) # Commands that trigger compilation/reload shouldn't retry on disconnect @@ -952,7 +958,8 @@ async def async_send_command_with_retry( loop=None, max_retries: int | None = None, retry_ms: int | None = None, - retry_on_reload: bool = True + retry_on_reload: bool = True, + budget_s: float | None = None, ) -> dict[str, Any] | MCPResponse: """Async wrapper that runs the blocking retry helper in a thread pool. @@ -964,6 +971,7 @@ async def async_send_command_with_retry( max_retries: Maximum number of retries for reload states retry_ms: Delay between retries in milliseconds retry_on_reload: If False, don't retry when Unity is reloading + budget_s: Optional cap in seconds on the whole call (see send_command_with_retry) Returns: Response dictionary or MCPResponse on error @@ -976,7 +984,7 @@ async def async_send_command_with_retry( None, lambda: send_command_with_retry( command_type, params, instance_id=instance_id, max_retries=max_retries, - retry_ms=retry_ms, retry_on_reload=retry_on_reload), + retry_ms=retry_ms, retry_on_reload=retry_on_reload, budget_s=budget_s), ) # After a successful command, check if the connection was freshly diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index 808b6ff153..fa1eb89749 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -166,3 +166,63 @@ async def hang(_send_fn, _instance, _command, _params, **kwargs): # noqa: ARG00 assert result.success is False assert "Timeout waiting for long_running_tool" in result.message assert elapsed < 2.0 + + +@pytest.mark.asyncio +async def test_a_stdio_poll_stops_its_worker_thread_at_the_polling_deadline(monkeypatch): + """On stdio the poll runs in an executor thread that a cancelled wait cannot stop. It must be given the + remaining budget, so the thread itself ends at the polling deadline and releases the connection.""" + import time + + import transport.legacy.unity_connection as unity_connection + + calls: list[dict] = [] + + class _BlockingConnection: + def send_command(self, command_type, params, max_attempts=None, deadline=None): # noqa: ARG002 + call = {"deadline": deadline, "ended": None} + calls.append(call) + started = time.monotonic() + # Block like a wedged socket would, until the deadline the caller allows (5 s at most here). + time.sleep(min(5.0, max(0.0, (deadline or started + 5.0) - started))) + call["ended"] = time.monotonic() + return {"_mcp_status": "pending", "_mcp_poll_interval": 0.01} + + monkeypatch.setattr(config, "transport_mode", "stdio") + monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None: _BlockingConnection()) + + service = _make_service() + started = time.monotonic() + result = await service._poll_until_complete( + "long_running_tool", None, {}, + {"_mcp_status": "pending", "_mcp_poll_interval": 0.01}, "status", max_poll_seconds=1) + await asyncio.sleep(0.2) # let a thread still blocked past the deadline (the bug) record itself + + assert result.success is False + assert calls, "the poll must have reached the stdio connection" + assert all(c["deadline"] is not None and c["deadline"] <= started + 1.5 for c in calls) + assert all(c["ended"] is not None and c["ended"] <= started + 1.5 for c in calls), "a poll thread outlived the deadline" + + +def test_budget_only_shortens_the_configured_command_deadline(monkeypatch): + import time + + import transport.legacy.unity_connection as unity_connection + + deadlines: list[float] = [] + + class _Connection: + def send_command(self, command_type, params, max_attempts=None, deadline=None): # noqa: ARG002 + deadlines.append(deadline - time.monotonic()) + return {"success": True} + + monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None: _Connection()) + monkeypatch.setattr(config, "command_total_timeout", 90.0) + + unity_connection.send_command_with_retry("x", {}, budget_s=2.0) + unity_connection.send_command_with_retry("x", {}, budget_s=500.0) + unity_connection.send_command_with_retry("x", {}) + + assert deadlines[0] <= 2.0 + assert 85.0 < deadlines[1] <= 90.0 + assert 85.0 < deadlines[2] <= 90.0 From 30c75a26333b50edca78bb479b7a522b94f36810 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Thu, 29 Oct 2026 23:02:48 +0900 Subject: [PATCH 6/9] fix(transport): the stdio budget starts before the connection lookup send_command_with_retry created its deadline after get_unity_connection(), so the time spent finding the instance, connecting and handshaking did not count, and a command whose caller had already given up (a status poll past its deadline) could still be sent. The deadline is now created first and passed through get_unity_connection, UnityConnectionPool.get_connection and UnityConnection.connect, which caps its connect and handshake timeouts to it (all optional arguments); if the lookup spent the whole budget, the command is not sent. Tests: a slow lookup leaves less of the budget, and one that spends it all sends nothing; connect caps its socket timeout to the deadline (both fail before). The existing deadline test's get_unity_connection fake accepts the new argument. --- .../src/transport/legacy/unity_connection.py | 43 +++++++++----- .../integration/test_connection_deadline.py | 2 +- Server/tests/test_custom_tool_timeouts.py | 57 ++++++++++++++++++- 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index 1440fa50d4..e42c9f0290 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -50,8 +50,11 @@ def _prepare_socket(self, sock: socket.socket) -> None: except OSError as exc: logger.debug(f"Unable to set TCP_NODELAY: {exc}") - def connect(self, connect_timeout: float | None = None) -> bool: - """Establish a connection to the Unity Editor.""" + def connect(self, connect_timeout: float | None = None, deadline: float | None = None) -> bool: + """Establish a connection to the Unity Editor. + + ``deadline`` (a monotonic() ceiling) caps both the connect and the handshake timeouts. + """ if self.sock: return True with self._conn_lock: @@ -62,6 +65,7 @@ def connect(self, connect_timeout: float | None = None) -> bool: if connect_timeout is None: connect_timeout = float( getattr(config, "connection_timeout", 1.0)) + connect_timeout = self._cap_to_deadline(connect_timeout, deadline) # 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( @@ -73,8 +77,8 @@ def connect(self, connect_timeout: float | None = None) -> bool: # Strict handshake: require FRAMING=1 try: require_framing = getattr(config, "require_framing", True) - handshake_timeout = float( - getattr(config, "handshake_timeout", 1.0)) + handshake_timeout = self._cap_to_deadline(float( + getattr(config, "handshake_timeout", 1.0)), deadline) self.sock.settimeout(handshake_timeout) buf = bytearray() deadline = time.monotonic() + handshake_timeout @@ -670,13 +674,14 @@ def _resolve_instance_id(self, instance_identifier: str | None, instances: list[ f"Check mcpforunity://instances resource for all instances." ) - def get_connection(self, instance_identifier: str | None = None) -> UnityConnection: + def get_connection(self, instance_identifier: str | None = None, deadline: float | None = None) -> UnityConnection: """ Get or create a connection to a Unity instance. Args: instance_identifier: Optional identifier (name, hash, name@hash, etc.) If None, uses default or most recent instance + deadline: Optional monotonic() ceiling for establishing a new connection Returns: UnityConnection to the specified instance @@ -696,7 +701,7 @@ def get_connection(self, instance_identifier: str | None = None) -> UnityConnect logger.info( f"Creating new connection to Unity instance: {target.id} (port {target.port})") conn = UnityConnection(port=target.port, instance_id=target.id) - if not conn.connect(): + if not conn.connect(deadline=deadline): raise ConnectionError( f"Failed to connect to Unity instance '{target.id}' on port {target.port}. " f"Ensure the Unity Editor is running." @@ -749,7 +754,7 @@ def get_unity_connection_pool() -> UnityConnectionPool: # Backwards compatibility: keep old single-connection function -def get_unity_connection(instance_identifier: str | None = None) -> UnityConnection: +def get_unity_connection(instance_identifier: str | None = None, deadline: float | None = None) -> UnityConnection: """Retrieve or establish a Unity connection. Args: @@ -762,7 +767,7 @@ def get_unity_connection(instance_identifier: str | None = None) -> UnityConnect Note: This function now uses the connection pool internally. """ pool = get_unity_connection_pool() - return pool.get_connection(instance_identifier) + return pool.get_connection(instance_identifier, deadline=deadline) # ----------------------------- @@ -842,9 +847,23 @@ def send_command_with_retry( """ t_retry_start = time.time() logger.info("[TIMING-STDIO] send_command_with_retry START command=%s", command_type) + # The budget starts now, so finding and connecting to the instance spends it too. + 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 budget_s is not None: + capped = time.monotonic() + max(0.0, budget_s) + deadline = capped if deadline is None else min(deadline, capped) t_get_conn = time.time() - conn = get_unity_connection(instance_id) + conn = get_unity_connection(instance_id, deadline=deadline) logger.info("[TIMING-STDIO] get_unity_connection took %.3fs command=%s", time.time() - t_get_conn, command_type) + if deadline is not None and time.monotonic() >= deadline: + # The caller has given up by now (a status poll past its deadline): do not send a command it no + # longer waits for. + return MCPResponse( + success=False, + error=f"No time left to send '{command_type}' after connecting to Unity", + hint="retry", + ) if max_retries is None: max_retries = getattr(config, "reload_max_retries", 40) if retry_ms is None: @@ -871,12 +890,6 @@ 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 budget_s is not None: - capped = time.monotonic() + max(0.0, budget_s) - deadline = capped if deadline is None else min(deadline, capped) - # 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 diff --git a/Server/tests/integration/test_connection_deadline.py b/Server/tests/integration/test_connection_deadline.py index dd6b280e43..5f1ca97c8c 100644 --- a/Server/tests/integration/test_connection_deadline.py +++ b/Server/tests/integration/test_connection_deadline.py @@ -91,7 +91,7 @@ def test_send_command_against_wedged_socket_is_bounded(silent_bridge, monkeypatc 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) + monkeypatch.setattr(uc, "get_unity_connection", lambda instance_id=None, deadline=None: conn) _write_reloading_status(tmp_path) start = time.monotonic() diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index fa1eb89749..a374adc83a 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -189,7 +189,7 @@ def send_command(self, command_type, params, max_attempts=None, deadline=None): return {"_mcp_status": "pending", "_mcp_poll_interval": 0.01} monkeypatch.setattr(config, "transport_mode", "stdio") - monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None: _BlockingConnection()) + monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None, deadline=None: _BlockingConnection()) service = _make_service() started = time.monotonic() @@ -216,7 +216,7 @@ def send_command(self, command_type, params, max_attempts=None, deadline=None): deadlines.append(deadline - time.monotonic()) return {"success": True} - monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None: _Connection()) + monkeypatch.setattr(unity_connection, "get_unity_connection", lambda instance_id=None, deadline=None: _Connection()) monkeypatch.setattr(config, "command_total_timeout", 90.0) unity_connection.send_command_with_retry("x", {}, budget_s=2.0) @@ -226,3 +226,56 @@ def send_command(self, command_type, params, max_attempts=None, deadline=None): assert deadlines[0] <= 2.0 assert 85.0 < deadlines[1] <= 90.0 assert 85.0 < deadlines[2] <= 90.0 + + +def test_finding_the_instance_spends_the_budget(monkeypatch): + """The budget starts before the connection lookup: a slow lookup leaves less for the command, and a + lookup that spends all of it means the command is not sent at all.""" + import time + + import transport.legacy.unity_connection as unity_connection + + sent: list[float] = [] + lookup_deadlines: list[float | None] = [] + + class _Connection: + def send_command(self, command_type, params, max_attempts=None, deadline=None): # noqa: ARG002 + sent.append(deadline - time.monotonic()) + return {"success": True} + + def slow_lookup(instance_id=None, deadline=None): # noqa: ARG001 + lookup_deadlines.append(deadline) + time.sleep(0.3) + return _Connection() + + monkeypatch.setattr(unity_connection, "get_unity_connection", slow_lookup) + monkeypatch.setattr(config, "command_total_timeout", 90.0) + + started = time.monotonic() + unity_connection.send_command_with_retry("x", {}, budget_s=1.0) + assert lookup_deadlines[0] is not None and lookup_deadlines[0] <= started + 1.05 + assert sent and sent[0] <= 0.75 + + result = unity_connection.send_command_with_retry("x", {}, budget_s=0.2) + assert len(sent) == 1, "a command whose budget the lookup spent must not be sent" + assert result.success is False + + +def test_connect_caps_its_timeouts_to_the_deadline(monkeypatch): + import socket + import time + + import transport.legacy.unity_connection as unity_connection + + timeouts: list[float] = [] + + def refuse(address, timeout=None): # noqa: ARG001 + timeouts.append(timeout) + raise ConnectionRefusedError("no editor here") + + monkeypatch.setattr(socket, "create_connection", refuse) + monkeypatch.setattr(config, "connection_timeout", 30.0) + + conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") + assert conn.connect(deadline=time.monotonic() + 0.5) is False + assert timeouts and timeouts[0] <= 0.5 From 8307d6795302db440c7a9b1272ca0ca419d769c1 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Fri, 30 Oct 2026 11:18:29 +0900 Subject: [PATCH 7/9] fix(transport): a stdio command never outlives its deadline on the socket - The I/O lock is awaited only for what is left of the budget, and nothing is sent once the deadline has passed; a command that times out waiting leaves the socket to the command holding it. - Sends and every receive are capped to the remaining budget, so a response arriving in pieces, or a payload Unity does not read, cannot hold the connection past it. - The liveness peek restores the socket timeout instead of clearing it with setblocking(True). - A reconnect passes the deadline to the handshake. - A custom tool's unity_timeout carries no retry hint: it may have run, so its effects are to be checked rather than the tool resubmitted. --- Server/src/services/custom_tool_service.py | 11 +- .../src/transport/legacy/unity_connection.py | 108 +++++++++----- .../integration/test_connection_deadline.py | 2 +- Server/tests/test_custom_tool_timeouts.py | 133 ++++++++++++++++++ 4 files changed, 213 insertions(+), 41 deletions(-) diff --git a/Server/src/services/custom_tool_service.py b/Server/src/services/custom_tool_service.py index 8c15152f93..e5f96dcadd 100644 --- a/Server/src/services/custom_tool_service.py +++ b/Server/src/services/custom_tool_service.py @@ -308,13 +308,18 @@ def _normalize_response(self, response) -> MCPResponse: if isinstance(response, MCPResponse): return response if isinstance(response, dict): + data = response.get( + "data", response) if "data" not in response else response["data"] + hint = response.get("hint") + if isinstance(data, dict) and data.get("reason") == "unity_timeout": + # The tool may have run: its effects are to be checked, not the tool resubmitted. + hint = None return MCPResponse( success=response.get("success", True), message=response.get("message"), error=response.get("error"), - data=response.get( - "data", response) if "data" not in response else response["data"], - hint=response.get("hint"), + data=data, + hint=hint, ) success = True diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index e42c9f0290..ef9cf36959 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -27,6 +27,10 @@ FRAMED_MAX = 64 * 1024 * 1024 +class _DeadlineBeforeSend(TimeoutError): + """The command's deadline passed before its payload was written: nothing was sent, the socket is intact.""" + + @dataclass class UnityConnection: """Manages the socket connection to the Unity Editor.""" @@ -144,9 +148,13 @@ def _ensure_live_connection(self) -> None: """ if not self.sock: return - orig_blocking = None + orig_timeout = None + peeked = False try: - orig_blocking = self.sock.getblocking() + # setblocking(True) would clear the timeout, leaving the next send and receive unbounded: + # restore the timeout itself instead. + orig_timeout = self.sock.gettimeout() + peeked = True self.sock.setblocking(False) data = self.sock.recv(1, socket.MSG_PEEK) if not data: @@ -161,21 +169,35 @@ def _ensure_live_connection(self) -> None: pass self.sock = None finally: - if self.sock and orig_blocking is not None: - self.sock.setblocking(orig_blocking) - - def _read_exact(self, sock: socket.socket, count: int) -> bytes: + if self.sock and peeked: + self.sock.settimeout(orig_timeout) + + def _recv_by(self, sock: socket.socket, size: int, deadline: float | None, timeout: float | None) -> bytes: + """One recv that waits no longer than ``timeout`` nor past ``deadline`` (a monotonic() ceiling).""" + if deadline is not None: + if time.monotonic() >= deadline: + raise socket.timeout("Deadline reached while receiving") + sock.settimeout(self._cap_to_deadline(timeout or config.connection_timeout, deadline)) + return sock.recv(size) + + def _read_exact(self, sock: socket.socket, count: int, deadline: float | None = None, + timeout: float | None = None) -> bytes: data = bytearray() while len(data) < count: - chunk = sock.recv(count - len(data)) + chunk = self._recv_by(sock, count - len(data), deadline, timeout) if not chunk: raise ConnectionError( "Connection closed before reading expected bytes") data.extend(chunk) return bytes(data) - def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: - """Receive a complete response from Unity, handling chunked data.""" + def receive_full_response(self, sock, buffer_size=config.buffer_size, deadline: float | None = None) -> bytes: + """Receive a complete response from Unity, handling chunked data. + + With a ``deadline`` (a monotonic() ceiling), every read is capped to what is left of it, so a + response arriving in pieces cannot hold the connection past the deadline. + """ + timeout = sock.gettimeout() if self.use_framing: # Heartbeat semantics: the Unity editor emits zero-length frames while # a long-running command is still executing. We tolerate a bounded @@ -187,7 +209,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: heartbeat_count = 0 try: while True: - header = self._read_exact(sock, 8) + header = self._read_exact(sock, 8, deadline, timeout) payload_len = struct.unpack('>Q', header)[0] if payload_len == 0: heartbeat_count += 1 @@ -201,7 +223,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: if payload_len > FRAMED_MAX: raise ValueError( f"Invalid framed length: {payload_len}") - payload = self._read_exact(sock, payload_len) + payload = self._read_exact(sock, payload_len, deadline, timeout) logger.debug( f"Received framed response ({len(payload)} bytes)") return payload @@ -218,7 +240,7 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: # Respect the socket's currently configured timeout try: while True: - chunk = sock.recv(buffer_size) + chunk = self._recv_by(sock, buffer_size, deadline, timeout) if not chunk: if not chunks: raise Exception( @@ -369,7 +391,7 @@ def read_status_file(target_hash: str | None = None) -> dict | None: self._ensure_live_connection() # Ensure connected (handshake occurs within connect()) t_conn_start = time.time() - if not self.sock and not self.connect(self._cap_to_deadline(config.connection_timeout, deadline)): + if not self.sock and not self.connect(self._cap_to_deadline(config.connection_timeout, deadline), deadline=deadline): raise ConnectionError("Could not connect to Unity") logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type) @@ -382,42 +404,51 @@ def read_status_file(target_hash: str | None = None) -> dict | None: 'params': params, }).encode('utf-8') - # Send/receive are serialized to protect the shared socket - with self._io_lock: + # Send/receive are serialized to protect the shared socket. Wait for another command's + # turn only as long as this one's budget lasts, and do not send once it has run out. + lock_wait = -1 if deadline is None else max(0.0, deadline - time.monotonic()) + if not self._io_lock.acquire(timeout=lock_wait): + raise _DeadlineBeforeSend( + f"Command '{command_type}' ran out of time waiting for the connection to Unity") + try: + if deadline is not None and time.monotonic() >= deadline: + raise _DeadlineBeforeSend( + f"Command '{command_type}' ran out of time before it was sent") mode = 'framed' if self.use_framing else 'legacy' with contextlib.suppress(Exception): logger.debug( f"send {len(payload)} bytes; mode={mode}; head={payload[:32].decode('utf-8', 'ignore')}") - t_send_start = time.time() - if self.use_framing: - header = struct.pack('>Q', len(payload)) - self.sock.sendall(header) - self.sock.sendall(payload) - else: - self.sock.sendall(payload) - logger.info("[TIMING-STDIO] sendall took %.3fs command=%s", time.time() - t_send_start, command_type) - - # 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 - recv_timeout = 1.0 if attempt > 0 else self.sock.gettimeout() + # Cap sends and receives to the remaining command budget (and use a short timeout + # during retry bursts) so a wedged socket can't block past the deadline. + original_timeout = self.sock.gettimeout() + io_timeout = 1.0 if attempt > 0 else original_timeout 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(recv_timeout) + io_timeout = self._cap_to_deadline( + io_timeout or config.connection_timeout, deadline) try: + self.sock.settimeout(io_timeout) + t_send_start = time.time() + if self.use_framing: + header = struct.pack('>Q', len(payload)) + self.sock.sendall(header) + self.sock.sendall(payload) + else: + self.sock.sendall(payload) + logger.info("[TIMING-STDIO] sendall took %.3fs command=%s", time.time() - t_send_start, command_type) + + if deadline is not None: + self.sock.settimeout(self._cap_to_deadline(io_timeout, deadline)) t_recv_start = time.time() - response_data = self.receive_full_response(self.sock) + response_data = self.receive_full_response(self.sock, deadline=deadline) logger.info("[TIMING-STDIO] receive took %.3fs command=%s len=%d", time.time() - t_recv_start, command_type, len(response_data)) with contextlib.suppress(Exception): logger.debug( f"recv {len(response_data)} bytes; mode={mode}") finally: - if restore_timeout is not None: - self.sock.settimeout(restore_timeout) + if self.sock: + self.sock.settimeout(original_timeout) + finally: + self._io_lock.release() # Parse if command_type == 'ping': @@ -432,6 +463,9 @@ def read_status_file(target_hash: str | None = None) -> dict | None: 'message', 'Unknown Unity error') raise Exception(err) return resp.get('result', {}) + except _DeadlineBeforeSend: + # Nothing was sent, and the socket may be serving another command: leave it alone. + raise except Exception as e: logger.warning( f"Unity communication attempt {attempt+1} failed: {e}") diff --git a/Server/tests/integration/test_connection_deadline.py b/Server/tests/integration/test_connection_deadline.py index 5f1ca97c8c..cfc1ff76f1 100644 --- a/Server/tests/integration/test_connection_deadline.py +++ b/Server/tests/integration/test_connection_deadline.py @@ -143,7 +143,7 @@ def discover_instances() -> list[UnityInstanceInfo]: attempted_ports: list[int] = [] ports_before_backoff: list[int] = [] - def fail_connect(connect_timeout: float | None = None) -> bool: + def fail_connect(connect_timeout: float | None = None, deadline: float | None = None) -> bool: attempted_ports.append(conn.port) return False diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index a374adc83a..5c03dd39ec 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -279,3 +279,136 @@ def refuse(address, timeout=None): # noqa: ARG001 conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") assert conn.connect(deadline=time.monotonic() + 0.5) is False assert timeouts and timeouts[0] <= 0.5 + + +def test_a_timed_out_custom_tool_carries_no_retry_hint(): + """A timeout may follow a command that ran: the caller is told to check, not to resubmit.""" + service = _make_service() + result = service._normalize_response( + {"success": False, "error": "late", "hint": "retry", "data": {"reason": "unity_timeout"}}) + assert result.hint is None + assert result.data == {"reason": "unity_timeout"} + + +def _call_in_thread(fn, limit_s): + """Run fn in a thread; return (elapsed, outcome) or fail if it is still blocked after limit_s.""" + import threading + import time + + outcome: dict = {} + + def run(): + try: + outcome["result"] = fn() + except Exception as exc: # noqa: BLE001 + outcome["error"] = exc + + started = time.monotonic() + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(limit_s) + assert not worker.is_alive(), f"still blocked after {limit_s} s" + return time.monotonic() - started, outcome + + +def _stdio_connection(monkeypatch, tmp_path): + import socket + + import transport.legacy.unity_connection as unity_connection + + # No status files: the reload preflight must not see a real editor's. + monkeypatch.setattr(unity_connection.Path, "home", lambda: tmp_path) + monkeypatch.setattr(unity_connection.stdio_port_registry, "get_instances", lambda **_: []) + monkeypatch.setattr(unity_connection.stdio_port_registry, "get_instance", lambda *_: None) + monkeypatch.setattr(unity_connection.stdio_port_registry, "get_port", lambda *_: 1) + ours, unity = socket.socketpair() + ours.settimeout(30.0) + conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") + conn.sock = ours + conn.use_framing = True + return conn, unity + + +def test_a_command_waits_for_the_connection_only_while_its_budget_lasts(monkeypatch, tmp_path): + import time + + conn, unity = _stdio_connection(monkeypatch, tmp_path) + unity.setblocking(False) + conn._io_lock.acquire() # another command is using the connection + try: + elapsed, outcome = _call_in_thread( + lambda: conn.send_command("x", {}, max_attempts=0, deadline=time.monotonic() + 0.3), 3.0) + finally: + conn._io_lock.release() + assert isinstance(outcome.get("error"), TimeoutError) + assert elapsed < 1.0 + assert conn.sock is not None, "the other command's socket must be left alone" + try: + sent = unity.recv(64) + except BlockingIOError: + sent = b"" + assert sent == b"", "a command whose budget ran out waiting must not be sent" + + +def test_a_response_in_pieces_is_cut_off_at_the_deadline(monkeypatch, tmp_path): + import struct + import threading + import time + + conn, unity = _stdio_connection(monkeypatch, tmp_path) + stop = threading.Event() + + def trickle(): + unity.sendall(struct.pack(">Q", 1000)) + while not stop.is_set(): # one byte every 0.2 s, never the whole payload + try: + unity.sendall(b"x") + except OSError: + return + time.sleep(0.2) + + threading.Thread(target=trickle, daemon=True).start() + try: + elapsed, outcome = _call_in_thread( + lambda: conn.receive_full_response(conn.sock, deadline=time.monotonic() + 0.5), 5.0) + finally: + stop.set() + assert isinstance(outcome.get("error"), TimeoutError) + assert elapsed < 1.0 + + +def test_a_send_unity_does_not_read_is_cut_off_at_the_deadline(monkeypatch, tmp_path): + import time + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + big = {"blob": "x" * (8 * 1024 * 1024)} # far beyond the socket buffers; Unity reads none of it + elapsed, outcome = _call_in_thread( + lambda: conn.send_command("x", big, max_attempts=0, deadline=time.monotonic() + 0.5), 5.0) + assert isinstance(outcome.get("error"), TimeoutError) + assert elapsed < 1.5 + + +def test_the_liveness_check_keeps_the_socket_timeout(monkeypatch, tmp_path): + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + conn._ensure_live_connection() + assert conn.sock is not None + assert conn.sock.gettimeout() == 30.0 + + +def test_a_reconnect_handshake_spends_the_command_budget(monkeypatch, tmp_path): + import time + + import transport.legacy.unity_connection as unity_connection + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + conn.sock = None # the stale socket was dropped: send_command reconnects + seen: list = [] + + def fake_connect(self, connect_timeout=None, deadline=None): # noqa: ARG001 + seen.append(deadline) + return False + + monkeypatch.setattr(unity_connection.UnityConnection, "connect", fake_connect) + deadline = time.monotonic() + 0.5 + _call_in_thread(lambda: conn.send_command("x", {}, max_attempts=0, deadline=deadline), 3.0) + assert seen == [deadline] From 36f5fab30c07abb4f36b3cb5a96c30d3a1443ba9 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Fri, 30 Oct 2026 11:37:56 +0900 Subject: [PATCH 8/9] fix(transport): the reload preflight waits for the connection only while the budget lasts When Unity reports a reload, send_command drops the socket under the I/O lock. That lock is now awaited only until the command's deadline; if another command still holds it, that command owns the socket and is left alone. --- .../src/transport/legacy/unity_connection.py | 16 ++++++++++++---- Server/tests/test_custom_tool_timeouts.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index ef9cf36959..a86a4180a1 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -293,6 +293,10 @@ def receive_full_response(self, sock, buffer_size=config.buffer_size, deadline: logger.error(f"Error during receive: {str(e)}") raise + def _lock_wait(self, deadline: float | None) -> float: + """How long to wait for the I/O lock: until the deadline (-1, without limit, when there is none).""" + return -1 if deadline is None else max(0.0, deadline - time.monotonic()) + 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: @@ -367,8 +371,13 @@ def read_status_file(target_hash: str | None = None) -> dict | None: 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() + # Wait for the lock only while the budget lasts: if another command still + # holds it, that command owns the socket and the reload ends its exchange. + if self._io_lock.acquire(timeout=self._lock_wait(deadline)): + try: + self.disconnect() + finally: + self._io_lock.release() return MCPResponse( success=False, error="Unity is reloading; please retry", @@ -406,8 +415,7 @@ def read_status_file(target_hash: str | None = None) -> dict | None: # Send/receive are serialized to protect the shared socket. Wait for another command's # turn only as long as this one's budget lasts, and do not send once it has run out. - lock_wait = -1 if deadline is None else max(0.0, deadline - time.monotonic()) - if not self._io_lock.acquire(timeout=lock_wait): + if not self._io_lock.acquire(timeout=self._lock_wait(deadline)): raise _DeadlineBeforeSend( f"Command '{command_type}' ran out of time waiting for the connection to Unity") try: diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index 5c03dd39ec..9958615496 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -412,3 +412,22 @@ def fake_connect(self, connect_timeout=None, deadline=None): # noqa: ARG001 deadline = time.monotonic() + 0.5 _call_in_thread(lambda: conn.send_command("x", {}, max_attempts=0, deadline=deadline), 3.0) assert seen == [deadline] + + +def test_the_reload_preflight_waits_for_the_connection_only_while_the_budget_lasts(monkeypatch, tmp_path): + import json + import time + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + status_dir = tmp_path / ".unity-mcp" + status_dir.mkdir() + (status_dir / "unity-mcp-status-abc.json").write_text(json.dumps({"reloading": True})) + conn._io_lock.acquire() # another command is using the connection + try: + elapsed, outcome = _call_in_thread( + lambda: conn.send_command("x", {}, max_attempts=0, deadline=time.monotonic() + 0.3), 3.0) + finally: + conn._io_lock.release() + assert elapsed < 1.0 + assert outcome["result"].hint == "retry" + assert conn.sock is not None, "the socket of the command holding the lock is left alone" From e33a9727f9c35add848d5066c57548acd8ba9009 Mon Sep 17 00:00:00 2001 From: Marco Meyer-Conde Date: Fri, 30 Oct 2026 11:56:10 +0900 Subject: [PATCH 9/9] fix(transport): serialize the whole exchange on the connection; keep both handshake limits - send_command takes the I/O lock before the liveness check and holds it through connect, send, receive and the close after a failure, so no command changes the socket or its timeout while another is using it. Port rediscovery and the backoff sleep run after the lock is released. - The handshake ends at the earlier of the configured handshake timeout and the caller's deadline, and each of its reads is capped to that, instead of rebasing the caller's deadline. --- .../src/transport/legacy/unity_connection.py | 233 +++++++++--------- Server/tests/test_custom_tool_timeouts.py | 46 ++++ 2 files changed, 166 insertions(+), 113 deletions(-) diff --git a/Server/src/transport/legacy/unity_connection.py b/Server/src/transport/legacy/unity_connection.py index a86a4180a1..de42d716e7 100644 --- a/Server/src/transport/legacy/unity_connection.py +++ b/Server/src/transport/legacy/unity_connection.py @@ -85,10 +85,13 @@ def connect(self, connect_timeout: float | None = None, deadline: float | None = getattr(config, "handshake_timeout", 1.0)), deadline) self.sock.settimeout(handshake_timeout) buf = bytearray() - deadline = time.monotonic() + handshake_timeout - while time.monotonic() < deadline and len(buf) < 512: + # The earlier of the configured handshake timeout and the caller's deadline. + handshake_deadline = time.monotonic() + handshake_timeout + if deadline is not None: + handshake_deadline = min(deadline, handshake_deadline) + while time.monotonic() < handshake_deadline and len(buf) < 512: try: - chunk = self.sock.recv(256) + chunk = self._recv_by(self.sock, 256, handshake_deadline, handshake_timeout) if not chunk: break buf.extend(chunk) @@ -394,31 +397,33 @@ def read_status_file(target_hash: str | None = None) -> dict | None: raise TimeoutError( f"Command '{command_type}' exceeded total deadline of " f"{total_timeout:.1f}s (connection wedged or Unity unresponsive)") + # Commands on this connection are serialized, from the liveness check to the close after a + # failure, so no command changes the socket (or its timeout) while another is using it. Wait for + # another command's turn only as long as this one's budget lasts. + if not self._io_lock.acquire(timeout=self._lock_wait(deadline)): + raise _DeadlineBeforeSend( + f"Command '{command_type}' ran out of time waiting for the connection to Unity") + failure: Exception | None = None 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(self._cap_to_deadline(config.connection_timeout, deadline), deadline=deadline): - raise ConnectionError("Could not connect to Unity") - logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type) - - # Build payload - if command_type == 'ping': - payload = b'ping' - else: - payload = json.dumps({ - 'type': command_type, - 'params': params, - }).encode('utf-8') - - # Send/receive are serialized to protect the shared socket. Wait for another command's - # turn only as long as this one's budget lasts, and do not send once it has run out. - if not self._io_lock.acquire(timeout=self._lock_wait(deadline)): - raise _DeadlineBeforeSend( - f"Command '{command_type}' ran out of time waiting for the connection to Unity") 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(self._cap_to_deadline(config.connection_timeout, deadline), deadline=deadline): + raise ConnectionError("Could not connect to Unity") + logger.info("[TIMING-STDIO] connect took %.3fs command=%s", time.time() - t_conn_start, command_type) + + # Build payload + if command_type == 'ping': + payload = b'ping' + else: + payload = json.dumps({ + 'type': command_type, + 'params': params, + }).encode('utf-8') + if deadline is not None and time.monotonic() >= deadline: raise _DeadlineBeforeSend( f"Command '{command_type}' ran out of time before it was sent") @@ -455,98 +460,100 @@ def read_status_file(target_hash: str | None = None) -> dict | None: finally: if self.sock: self.sock.settimeout(original_timeout) - finally: - self._io_lock.release() - # Parse - if command_type == 'ping': - resp = json.loads(response_data.decode('utf-8')) - if resp.get('status') == 'success' and resp.get('result', {}).get('message') == 'pong': - return {"message": "pong"} - raise Exception("Ping unsuccessful") - - resp = json.loads(response_data.decode('utf-8')) - if resp.get('status') == 'error': - err = resp.get('error') or resp.get( - 'message', 'Unknown Unity error') - raise Exception(err) - return resp.get('result', {}) - except _DeadlineBeforeSend: - # Nothing was sent, and the socket may be serving another command: leave it alone. - raise - except Exception as e: - logger.warning( - f"Unity communication attempt {attempt+1} failed: {e}") - try: - if self.sock: - self.sock.close() - finally: - self.sock = None + # Parse + if command_type == 'ping': + resp = json.loads(response_data.decode('utf-8')) + if resp.get('status') == 'success' and resp.get('result', {}).get('message') == 'pong': + return {"message": "pong"} + raise Exception("Ping unsuccessful") - # Re-discover the port for this specific instance - try: - stdio_port_registry.get_instances(force_refresh=True) - new_port: int | None = None - if self.instance_id: - # Try to rediscover the specific instance via shared registry - refreshed_instance = stdio_port_registry.get_instance( - self.instance_id) - if refreshed_instance and isinstance(refreshed_instance.port, int): - new_port = refreshed_instance.port - logger.debug( - f"Rediscovered instance {self.instance_id} on port {new_port}") - else: - logger.warning( - f"Instance {self.instance_id} not found during reconnection; falling back to port scan", - ) + resp = json.loads(response_data.decode('utf-8')) + if resp.get('status') == 'error': + err = resp.get('error') or resp.get( + 'message', 'Unknown Unity error') + raise Exception(err) + return resp.get('result', {}) + except _DeadlineBeforeSend: + # Nothing was sent, and the socket is left as it was. + raise + except Exception as e: + failure = e + logger.warning( + f"Unity communication attempt {attempt+1} failed: {e}") + try: + if self.sock: + self.sock.close() + finally: + self.sock = None + finally: + self._io_lock.release() - # Fallback to registry default if instance-specific discovery failed - if new_port is None: - new_port = stdio_port_registry.get_port( - self.instance_id) - logger.info( - f"Using Unity port from stdio_port_registry: {new_port}") - - if new_port != self.port: - logger.info( - f"Unity port changed {self.port} -> {new_port}") - self.port = new_port - except Exception as de: - logger.debug(f"Port discovery failed: {de}") - - if attempt < attempts: - # Heartbeat-aware, jittered backoff - status = read_status_file(target_hash) - # Base exponential backoff - backoff = base_backoff * (2 ** attempt) - # Decorrelated jitter multiplier - jitter = random.uniform(0.1, 0.3) - - # Fast‑retry for transient socket failures - fast_error = isinstance( - e, (ConnectionRefusedError, ConnectionResetError, TimeoutError)) - if not fast_error: - try: - err_no = getattr(e, 'errno', None) - fast_error = err_no in ( - errno.ECONNREFUSED, errno.ECONNRESET, errno.ETIMEDOUT) - except Exception: - pass - - # Cap backoff depending on state - if status and status.get('reloading'): - # Domain reload can take 10-20s; use longer waits - cap = 5.0 - elif fast_error: - cap = 0.25 + e = failure + # Re-discover the port for this specific instance + try: + stdio_port_registry.get_instances(force_refresh=True) + new_port: int | None = None + if self.instance_id: + # Try to rediscover the specific instance via shared registry + refreshed_instance = stdio_port_registry.get_instance( + self.instance_id) + if refreshed_instance and isinstance(refreshed_instance.port, int): + new_port = refreshed_instance.port + logger.debug( + f"Rediscovered instance {self.instance_id} on port {new_port}") else: - cap = 3.0 + logger.warning( + f"Instance {self.instance_id} not found during reconnection; falling back to port scan", + ) + + # Fallback to registry default if instance-specific discovery failed + if new_port is None: + new_port = stdio_port_registry.get_port( + self.instance_id) + logger.info( + f"Using Unity port from stdio_port_registry: {new_port}") - 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 + if new_port != self.port: + logger.info( + f"Unity port changed {self.port} -> {new_port}") + self.port = new_port + except Exception as de: + logger.debug(f"Port discovery failed: {de}") + + if attempt < attempts: + # Heartbeat-aware, jittered backoff + status = read_status_file(target_hash) + # Base exponential backoff + backoff = base_backoff * (2 ** attempt) + # Decorrelated jitter multiplier + jitter = random.uniform(0.1, 0.3) + + # Fast‑retry for transient socket failures + fast_error = isinstance( + e, (ConnectionRefusedError, ConnectionResetError, TimeoutError)) + if not fast_error: + try: + err_no = getattr(e, 'errno', None) + fast_error = err_no in ( + errno.ECONNREFUSED, errno.ECONNRESET, errno.ETIMEDOUT) + except Exception: + pass + + # Cap backoff depending on state + if status and status.get('reloading'): + # Domain reload can take 10-20s; use longer waits + cap = 5.0 + elif fast_error: + cap = 0.25 + else: + 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 e # ----------------------------- diff --git a/Server/tests/test_custom_tool_timeouts.py b/Server/tests/test_custom_tool_timeouts.py index 9958615496..1252e16040 100644 --- a/Server/tests/test_custom_tool_timeouts.py +++ b/Server/tests/test_custom_tool_timeouts.py @@ -431,3 +431,49 @@ def test_the_reload_preflight_waits_for_the_connection_only_while_the_budget_las assert elapsed < 1.0 assert outcome["result"].hint == "retry" assert conn.sock is not None, "the socket of the command holding the lock is left alone" + + +def test_a_waiting_command_leaves_the_socket_alone_until_its_turn(monkeypatch, tmp_path): + """The liveness check reads and resets the socket's timeout: done while another command holds the + connection, it could put back a longer timeout under that command's send.""" + import time + + conn, _unity = _stdio_connection(monkeypatch, tmp_path) + checked: list[bool] = [] + monkeypatch.setattr(conn, "_ensure_live_connection", lambda: checked.append(True)) + conn._io_lock.acquire() # another command is using the connection + try: + _call_in_thread( + lambda: conn.send_command("x", {}, max_attempts=0, deadline=time.monotonic() + 0.3), 3.0) + finally: + conn._io_lock.release() + assert checked == [], "the liveness check ran while another command held the connection" + + +def test_the_handshake_ends_at_the_earlier_of_its_timeout_and_the_deadline(monkeypatch): + import socket + import threading + import time + + import transport.legacy.unity_connection as unity_connection + + ours, unity = socket.socketpair() + stop = threading.Event() + + def trickle(): # a byte every 0.45 s, never the end of the handshake line + while not stop.is_set(): + try: + unity.sendall(b"F") + except OSError: + return + time.sleep(0.45) + + monkeypatch.setattr(socket, "create_connection", lambda *_a, **_k: ours) + monkeypatch.setattr(config, "handshake_timeout", 5.0) + threading.Thread(target=trickle, daemon=True).start() + conn = unity_connection.UnityConnection(port=1, instance_id="Playground@abc") + try: + elapsed, _ = _call_in_thread(lambda: conn.connect(deadline=time.monotonic() + 0.5), 5.0) + finally: + stop.set() + assert elapsed < 0.75