Skip to content

Commit 57e30b4

Browse files
committed
Prefer PluginHub for HTTP, limit stdio auto-select, and tighten exception handling
### Motivation - Avoid probing `stdio` when the server transport is `http` and prefer PluginHub session resolution instead. - Prevent transient errors from clearing or flipping the active instance by making auto-selection defensive. - Reduce noisy exception handling by catching expected errors explicitly and re-raising system-level interrupts. - Improve observability by logging successful auto-selection at `info` and probe failures at `debug`. ### Description - Added `async def _maybe_autoselect_instance(self, ctx)` which calls `transport.unity_transport._current_transport()` and tries PluginHub `get_sessions()` first, falling back to stdio discovery only when `transport != "http"`. - Limited stdio discovery to `get_unity_connection_pool().discover_all_instances(force_refresh=True)` and only auto-select when exactly one instance is found, and set the session via `set_active_instance`. - Replaced broad `except Exception:` catches with targeted exception tuples and a final handler that re-raises `SystemExit`/`KeyboardInterrupt`, and added clearer `info`/`debug` logging messages. - Updated `_inject_unity_instance` to call `_maybe_autoselect_instance` when there is no `active_instance` and to validate PluginHub session resolution defensively. ### Testing - No automated tests were executed as part of this rollout.
1 parent a09bb2a commit 57e30b4

File tree

3 files changed

+229
-0
lines changed

3 files changed

+229
-0
lines changed

Server/src/transport/unity_instance_middleware.py

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,11 +83,93 @@ def clear_active_instance(self, ctx) -> None:
8383
with self._lock:
8484
self._active_by_key.pop(key, None)
8585

86+
async def _maybe_autoselect_instance(self, ctx) -> str | None:
87+
"""Auto-select sole Unity instance when no active instance is set."""
88+
try:
89+
from transport.unity_transport import _current_transport
90+
91+
transport = _current_transport()
92+
if PluginHub.is_configured():
93+
try:
94+
sessions_data = await PluginHub.get_sessions()
95+
sessions = sessions_data.sessions or {}
96+
ids: list[str] = []
97+
for session_info in sessions.values():
98+
project = getattr(session_info, "project", None) or "Unknown"
99+
hash_value = getattr(session_info, "hash", None)
100+
if hash_value:
101+
ids.append(f"{project}@{hash_value}")
102+
if len(ids) == 1:
103+
chosen = ids[0]
104+
self.set_active_instance(ctx, chosen)
105+
logger.info(
106+
"Auto-selected sole Unity instance via PluginHub: %s",
107+
chosen,
108+
)
109+
return chosen
110+
except (ConnectionError, ValueError, KeyError, TimeoutError, AttributeError) as exc:
111+
logger.debug(
112+
"PluginHub auto-select probe failed (%s); falling back to stdio",
113+
type(exc).__name__,
114+
exc_info=True,
115+
)
116+
except Exception as exc:
117+
if isinstance(exc, (SystemExit, KeyboardInterrupt)):
118+
raise
119+
logger.debug(
120+
"PluginHub auto-select probe failed with unexpected error (%s); falling back to stdio",
121+
type(exc).__name__,
122+
exc_info=True,
123+
)
124+
125+
if transport != "http":
126+
try:
127+
from transport.legacy.unity_connection import get_unity_connection_pool
128+
129+
pool = get_unity_connection_pool()
130+
instances = pool.discover_all_instances(force_refresh=True)
131+
ids = [getattr(inst, "id", None) for inst in instances]
132+
ids = [inst_id for inst_id in ids if inst_id]
133+
if len(ids) == 1:
134+
chosen = ids[0]
135+
self.set_active_instance(ctx, chosen)
136+
logger.info(
137+
"Auto-selected sole Unity instance via stdio discovery: %s",
138+
chosen,
139+
)
140+
return chosen
141+
except (ConnectionError, ValueError, KeyError, TimeoutError, AttributeError) as exc:
142+
logger.debug(
143+
"Stdio auto-select probe failed (%s)",
144+
type(exc).__name__,
145+
exc_info=True,
146+
)
147+
except Exception as exc:
148+
if isinstance(exc, (SystemExit, KeyboardInterrupt)):
149+
raise
150+
logger.debug(
151+
"Stdio auto-select probe failed with unexpected error (%s)",
152+
type(exc).__name__,
153+
exc_info=True,
154+
)
155+
except Exception as exc:
156+
if isinstance(exc, (SystemExit, KeyboardInterrupt)):
157+
raise
158+
logger.debug(
159+
"Auto-select path encountered an unexpected error (%s)",
160+
type(exc).__name__,
161+
exc_info=True,
162+
)
163+
164+
return None
165+
86166
async def _inject_unity_instance(self, context: MiddlewareContext) -> None:
87167
"""Inject active Unity instance into context if available."""
88168
ctx = context.fastmcp_context
89169

90170
active_instance = self.get_active_instance(ctx)
171+
if not active_instance:
172+
active_instance = await self._maybe_autoselect_instance(ctx)
91173
if active_instance:
92174
# If using HTTP transport (PluginHub configured), validate session
93175
# But for stdio transport (no PluginHub needed or maybe partially configured),

Server/tests/integration/conftest.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
SERVER_ROOT = Path(__file__).resolve().parents[2]
77
if str(SERVER_ROOT) not in sys.path:
88
sys.path.insert(0, str(SERVER_ROOT))
9+
SERVER_SRC = SERVER_ROOT / "src"
10+
if str(SERVER_SRC) not in sys.path:
11+
sys.path.insert(0, str(SERVER_SRC))
912

1013
# Ensure telemetry is disabled during test collection and execution to avoid
1114
# any background network or thread startup that could slow or block pytest.
@@ -86,3 +89,39 @@ class _DummyMiddlewareContext:
8689
fastmcp_server.middleware = fastmcp_server_middleware
8790
sys.modules.setdefault("fastmcp.server", fastmcp_server)
8891
sys.modules.setdefault("fastmcp.server.middleware", fastmcp_server_middleware)
92+
93+
# Stub minimal starlette modules to avoid optional dependency imports.
94+
starlette = types.ModuleType("starlette")
95+
starlette_endpoints = types.ModuleType("starlette.endpoints")
96+
starlette_websockets = types.ModuleType("starlette.websockets")
97+
starlette_requests = types.ModuleType("starlette.requests")
98+
starlette_responses = types.ModuleType("starlette.responses")
99+
100+
101+
class _DummyWebSocketEndpoint:
102+
pass
103+
104+
105+
class _DummyWebSocket:
106+
pass
107+
108+
109+
class _DummyRequest:
110+
pass
111+
112+
113+
class _DummyJSONResponse:
114+
def __init__(self, *args, **kwargs):
115+
pass
116+
117+
118+
starlette_endpoints.WebSocketEndpoint = _DummyWebSocketEndpoint
119+
starlette_websockets.WebSocket = _DummyWebSocket
120+
starlette_requests.Request = _DummyRequest
121+
starlette_responses.JSONResponse = _DummyJSONResponse
122+
123+
sys.modules.setdefault("starlette", starlette)
124+
sys.modules.setdefault("starlette.endpoints", starlette_endpoints)
125+
sys.modules.setdefault("starlette.websockets", starlette_websockets)
126+
sys.modules.setdefault("starlette.requests", starlette_requests)
127+
sys.modules.setdefault("starlette.responses", starlette_responses)
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import asyncio
2+
import sys
3+
import types
4+
from types import SimpleNamespace
5+
6+
from .test_helpers import DummyContext
7+
8+
9+
class DummyMiddlewareContext:
10+
def __init__(self, ctx):
11+
self.fastmcp_context = ctx
12+
13+
14+
def test_auto_selects_single_instance_via_pluginhub(monkeypatch):
15+
plugin_hub = types.ModuleType("transport.plugin_hub")
16+
17+
class PluginHub:
18+
@classmethod
19+
def is_configured(cls) -> bool:
20+
return True
21+
22+
@classmethod
23+
async def get_sessions(cls):
24+
raise AssertionError("get_sessions should be stubbed in test")
25+
26+
plugin_hub.PluginHub = PluginHub
27+
monkeypatch.setitem(sys.modules, "transport.plugin_hub", plugin_hub)
28+
unity_transport = types.ModuleType("transport.unity_transport")
29+
unity_transport._current_transport = lambda: "http"
30+
monkeypatch.setitem(sys.modules, "transport.unity_transport", unity_transport)
31+
monkeypatch.delitem(sys.modules, "transport.unity_instance_middleware", raising=False)
32+
33+
from transport.unity_instance_middleware import UnityInstanceMiddleware, PluginHub as ImportedPluginHub
34+
assert ImportedPluginHub is plugin_hub.PluginHub
35+
36+
monkeypatch.setenv("UNITY_MCP_TRANSPORT", "http")
37+
38+
middleware = UnityInstanceMiddleware()
39+
ctx = DummyContext()
40+
ctx.client_id = "client-1"
41+
middleware_context = DummyMiddlewareContext(ctx)
42+
43+
call_count = {"sessions": 0}
44+
45+
async def fake_get_sessions():
46+
call_count["sessions"] += 1
47+
return SimpleNamespace(
48+
sessions={
49+
"session-1": SimpleNamespace(project="Ramble", hash="deadbeef"),
50+
}
51+
)
52+
53+
monkeypatch.setattr(plugin_hub.PluginHub, "get_sessions", fake_get_sessions)
54+
55+
selected = asyncio.run(middleware._maybe_autoselect_instance(ctx))
56+
57+
assert selected == "Ramble@deadbeef"
58+
assert middleware.get_active_instance(ctx) == "Ramble@deadbeef"
59+
assert call_count["sessions"] == 1
60+
61+
asyncio.run(middleware._inject_unity_instance(middleware_context))
62+
63+
assert ctx.get_state("unity_instance") == "Ramble@deadbeef"
64+
assert call_count["sessions"] == 1
65+
66+
67+
def test_auto_selects_single_instance_via_stdio(monkeypatch):
68+
plugin_hub = types.ModuleType("transport.plugin_hub")
69+
70+
class PluginHub:
71+
@classmethod
72+
def is_configured(cls) -> bool:
73+
return False
74+
75+
plugin_hub.PluginHub = PluginHub
76+
monkeypatch.setitem(sys.modules, "transport.plugin_hub", plugin_hub)
77+
unity_transport = types.ModuleType("transport.unity_transport")
78+
unity_transport._current_transport = lambda: "stdio"
79+
monkeypatch.setitem(sys.modules, "transport.unity_transport", unity_transport)
80+
monkeypatch.delitem(sys.modules, "transport.unity_instance_middleware", raising=False)
81+
82+
from transport.unity_instance_middleware import UnityInstanceMiddleware, PluginHub as ImportedPluginHub
83+
assert ImportedPluginHub is plugin_hub.PluginHub
84+
85+
monkeypatch.setenv("UNITY_MCP_TRANSPORT", "stdio")
86+
87+
middleware = UnityInstanceMiddleware()
88+
ctx = DummyContext()
89+
ctx.client_id = "client-1"
90+
middleware_context = DummyMiddlewareContext(ctx)
91+
92+
class PoolStub:
93+
def discover_all_instances(self, force_refresh=False):
94+
assert force_refresh is True
95+
return [SimpleNamespace(id="UnityMCPTests@cc8756d4")]
96+
97+
unity_connection = types.ModuleType("transport.legacy.unity_connection")
98+
unity_connection.get_unity_connection_pool = lambda: PoolStub()
99+
monkeypatch.setitem(sys.modules, "transport.legacy.unity_connection", unity_connection)
100+
101+
selected = asyncio.run(middleware._maybe_autoselect_instance(ctx))
102+
103+
assert selected == "UnityMCPTests@cc8756d4"
104+
assert middleware.get_active_instance(ctx) == "UnityMCPTests@cc8756d4"
105+
106+
asyncio.run(middleware._inject_unity_instance(middleware_context))
107+
108+
assert ctx.get_state("unity_instance") == "UnityMCPTests@cc8756d4"

0 commit comments

Comments
 (0)