diff --git a/Server/src/transport/unity_instance_middleware.py b/Server/src/transport/unity_instance_middleware.py index 82ba47fd7..238e01d54 100644 --- a/Server/src/transport/unity_instance_middleware.py +++ b/Server/src/transport/unity_instance_middleware.py @@ -5,6 +5,7 @@ into the request-scoped state, allowing tools to access it via ctx.get_state("unity_instance"). """ from threading import RLock +import asyncio import logging import time @@ -212,6 +213,57 @@ async def _resolve_instance_value(self, value: str, ctx) -> str: "Read mcpforunity://instances for current sessions." ) + async def _drop_stale_pin(self, ctx, active_instance: str) -> str: + """ + Discard a pin whose editor never came back, and refuse rather than retarget. + + A pin outlives the editor it names. Left in place it also suppresses + auto-select, so every later call fails with no_unity_session and neither + waiting nor relaunching the editor recovers: the editor re-registers + under its own name, which is not the name the pin holds. + + Two things make dropping it safe. The pinned hash gets the same reconnect + window `_resolve_session_id` gives it, so a domain reload does not look like + a departure. And a departure raises instead of returning None, because + falling through to auto-select would land a call pinned to project A in + project B, report success, and re-pin B (#1023). + + An empty registry means every editor is mid-reload, so the pin still stands. + """ + from transport.unity_transport import _is_http_transport + if not (_is_http_transport() and PluginHub.is_configured()): + return active_instance + + from transport.plugin_hub import _read_bounded_wait_env + max_wait_s = _read_bounded_wait_env( + "UNITY_MCP_SESSION_RESOLVE_MAX_WAIT_S", default_s=20.0, max_s=120.0) + retry_ms = float(getattr(config, "reload_retry_ms", 250)) + sleep_seconds = max(0.05, min(0.25, retry_ms / 1000.0)) + deadline = time.monotonic() + max_wait_s + + while True: + instances = await self._discover_instances(ctx) + ids = [inst.id for inst in instances if getattr(inst, "id", None)] + if not ids or active_instance in ids: + return active_instance + if time.monotonic() >= deadline: + break + await asyncio.sleep(sleep_seconds) + + _diag.warning( + "Active instance %s did not return within %.1fs; dropping the stale pin. Registered: %s", + active_instance, + max_wait_s, + ", ".join(ids), + ) + await self.clear_active_instance(ctx) + raise ValueError( + f"Pinned Unity instance '{active_instance}' is no longer running " + f"(waited {max_wait_s:.0f}s). The pin has been cleared. Available: " + f"{', '.join(ids) or 'none'}. Pass unity_instance explicitly, or read " + "mcpforunity://instances for current sessions." + ) + async def _maybe_autoselect_instance(self, ctx) -> str | None: """ Auto-select the sole Unity instance when no active instance is set. @@ -351,6 +403,8 @@ async def _inject_unity_instance(self, context: MiddlewareContext) -> None: if not active_instance: active_instance = await self.get_active_instance(ctx) + if active_instance: + active_instance = await self._drop_stale_pin(ctx, active_instance) if not active_instance: active_instance = await self._maybe_autoselect_instance(ctx) if active_instance: diff --git a/Server/tests/integration/test_instance_autoselect.py b/Server/tests/integration/test_instance_autoselect.py index b974cb11c..036e75ea0 100644 --- a/Server/tests/integration/test_instance_autoselect.py +++ b/Server/tests/integration/test_instance_autoselect.py @@ -26,6 +26,7 @@ async def get_sessions(cls): raise AssertionError("get_sessions should be stubbed in test") plugin_hub.PluginHub = PluginHub + plugin_hub._read_bounded_wait_env = lambda name, default_s, max_s: 0.0 monkeypatch.setitem(sys.modules, "transport.plugin_hub", plugin_hub) monkeypatch.delitem(sys.modules, "transport.unity_instance_middleware", raising=False) @@ -73,6 +74,7 @@ def is_configured(cls) -> bool: return False plugin_hub.PluginHub = PluginHub + plugin_hub._read_bounded_wait_env = lambda name, default_s, max_s: 0.0 monkeypatch.setitem(sys.modules, "transport.plugin_hub", plugin_hub) monkeypatch.delitem(sys.modules, "transport.unity_instance_middleware", raising=False) @@ -115,6 +117,7 @@ def is_configured(cls) -> bool: return False plugin_hub.PluginHub = PluginHub + plugin_hub._read_bounded_wait_env = lambda name, default_s, max_s: 0.0 monkeypatch.setitem(sys.modules, "transport.plugin_hub", plugin_hub) monkeypatch.delitem(sys.modules, "transport.unity_instance_middleware", raising=False) diff --git a/Server/tests/test_stale_instance_pin.py b/Server/tests/test_stale_instance_pin.py new file mode 100644 index 000000000..f9f1ac3d9 --- /dev/null +++ b/Server/tests/test_stale_instance_pin.py @@ -0,0 +1,136 @@ +"""Stale active-instance pin tests. + +A pin outlives the editor it names. While it stayed pinned it also suppressed +auto-select, so every later call failed with no_unity_session and neither +waiting nor relaunching the editor recovered: the editor re-registers under its +own name, which is not the name the pin holds. + +Dropping it is only safe under two conditions, both covered here: the pinned hash +gets the same reconnect window `_resolve_session_id` gives it, so a domain reload +does not read as a departure; and a departure raises rather than falling through to +auto-select, which would silently move a call from project A to project B (#1023). +""" + +from types import SimpleNamespace + +import pytest + +import transport.unity_transport as unity_transport +from transport.plugin_hub import PluginHub +from transport.unity_instance_middleware import UnityInstanceMiddleware + + +class FakeContext: + """Minimal ctx shim over the session-scoped state the middleware uses.""" + + def __init__(self, state: dict | None = None): + self._state = state or {} + + async def get_state(self, key: str): + return self._state.get(key) + + async def set_state(self, key: str, value) -> None: + self._state[key] = value + + +@pytest.fixture +def http_hub(monkeypatch): + monkeypatch.setattr(unity_transport, "_is_http_transport", lambda: True) + monkeypatch.setattr(PluginHub, "is_configured", classmethod(lambda cls: True)) + # No reconnect window unless a test asks for one, so the suite does not sit out 20s + monkeypatch.setenv("UNITY_MCP_SESSION_RESOLVE_MAX_WAIT_S", "0") + + +def _registered(monkeypatch, middleware, *ids): + async def discover(_ctx): + return [SimpleNamespace(id=i, hash=i.split("@")[-1], name=i.split("@")[0]) for i in ids] + + monkeypatch.setattr(middleware, "_discover_instances", discover) + + +def _registered_over_time(monkeypatch, middleware, *rounds): + """Return a different registry on each successive poll.""" + calls = {"n": 0} + + async def discover(_ctx): + ids = rounds[min(calls["n"], len(rounds) - 1)] + calls["n"] += 1 + return [SimpleNamespace(id=i, hash=i.split("@")[-1], name=i.split("@")[0]) for i in ids] + + monkeypatch.setattr(middleware, "_discover_instances", discover) + return calls + + +@pytest.mark.asyncio +async def test_departed_pin_is_dropped_and_refuses(http_hub, monkeypatch): + """The ghost pin that bricked routing must not survive a call.""" + middleware = UnityInstanceMiddleware() + ctx = FakeContext({UnityInstanceMiddleware._ACTIVE_INSTANCE_STATE_KEY: "wt-display-system@bbbb"}) + _registered(monkeypatch, middleware, "Trailblazers-1@aaaa") + + with pytest.raises(ValueError) as excinfo: + await middleware._drop_stale_pin(ctx, "wt-display-system@bbbb") + + assert "wt-display-system@bbbb" in str(excinfo.value) + assert await middleware.get_active_instance(ctx) is None + + +@pytest.mark.asyncio +async def test_pin_naming_a_registered_instance_is_kept(http_hub, monkeypatch): + """A live pin is the user's choice and must be left alone.""" + middleware = UnityInstanceMiddleware() + ctx = FakeContext() + _registered(monkeypatch, middleware, "Trailblazers-1@aaaa", "Trailblazers-2@cccc") + + resolved = await middleware._drop_stale_pin(ctx, "Trailblazers-2@cccc") + + assert resolved == "Trailblazers-2@cccc" + + +@pytest.mark.asyncio +async def test_pin_is_kept_while_no_instance_is_registered(http_hub, monkeypatch): + """An empty registry is a domain reload in flight, not a dead instance.""" + middleware = UnityInstanceMiddleware() + ctx = FakeContext() + _registered(monkeypatch, middleware) + + resolved = await middleware._drop_stale_pin(ctx, "Trailblazers-1@aaaa") + + assert resolved == "Trailblazers-1@aaaa" + + +@pytest.mark.asyncio +async def test_pin_survives_a_reload_that_ends_inside_the_wait(http_hub, monkeypatch): + """A reloading editor is absent from a populated registry; the wait must cover it.""" + monkeypatch.setenv("UNITY_MCP_SESSION_RESOLVE_MAX_WAIT_S", "5") + middleware = UnityInstanceMiddleware() + ctx = FakeContext() + calls = _registered_over_time( + monkeypatch, + middleware, + ("Trailblazers-1@aaaa",), + ("Trailblazers-1@aaaa", "Trailblazers-2@cccc"), + ) + + resolved = await middleware._drop_stale_pin(ctx, "Trailblazers-2@cccc") + + assert resolved == "Trailblazers-2@cccc" + assert calls["n"] > 1, "must poll again rather than drop on the first miss" + + +@pytest.mark.asyncio +async def test_dropped_pin_never_retargets_another_project(http_hub, monkeypatch): + """Clearing then auto-selecting the survivor is #1023 arriving through another door.""" + middleware = UnityInstanceMiddleware() + ctx = FakeContext({UnityInstanceMiddleware._ACTIVE_INSTANCE_STATE_KEY: "project-a@bbbb"}) + _registered(monkeypatch, middleware, "project-b@aaaa") + + async def autoselect(_ctx): + pytest.fail("auto-select must not run after a pinned instance departs") + + monkeypatch.setattr(middleware, "_maybe_autoselect_instance", autoselect) + + with pytest.raises(ValueError): + await middleware._drop_stale_pin(ctx, "project-a@bbbb") + + assert await middleware.get_active_instance(ctx) is None