fix(server): stop silent cross-instance routing when multiple editors are connected (#1023)

Server/src/services/tools/debug_request_context.py

@@ -46,16 +46,13 @@ async def debug_request_context(ctx: Context) -> dict[str, Any]:
     # List all ctx attributes for debugging
     ctx_attrs = [attr for attr in dir(ctx) if not attr.startswith("_")]
 
-    # Get session state info via middleware
+    # Get session state info via middleware. Active-instance storage now lives
+    # in FastMCP's session-scoped state store, keyed by ctx.session_id, so
+    # there is no global dict to enumerate — that snapshot was a footgun
+    # anyway (it exposed every connected client's selection).
     middleware = get_unity_instance_middleware()
-    derived_key = await middleware.get_session_key(ctx)
     active_instance = await middleware.get_active_instance(ctx)
 
-    # Debugging middleware internals
-    # NOTE: These fields expose internal implementation details and may change between versions.
-    with middleware._lock:
-        all_keys = list(middleware._active_by_key.keys())
-
     # Debugging PluginHub state
     plugin_hub_configured = PluginHub.is_configured()
 
@@ -77,9 +74,7 @@ async def debug_request_context(ctx: Context) -> dict[str, Any]:
                 "client_id": ctx_client_id,
             },
             "session_state": {
-                "derived_key": derived_key,
                 "active_instance": active_instance,
-                "all_keys_in_store": all_keys,
                 "plugin_hub_configured": plugin_hub_configured,
                 "middleware_id": id(middleware),
             },

Server/src/services/tools/set_active_instance.py

@@ -54,7 +54,7 @@ async def set_active_instance(
             "message": f"Active instance set to {resolved_id}",
             "data": {
                 "instance": resolved_id,
-                "session_key": await middleware.get_session_key(ctx),
+                "session_id": getattr(ctx, "session_id", None),
             },
         }
 
@@ -138,18 +138,15 @@ async def set_active_instance(
             "error": "Internal error: Instance resolution failed."
         }
 
-    # Store selection in middleware (session-scoped)
+    # Store selection in middleware (session-scoped via FastMCP context state).
     middleware = get_unity_instance_middleware()
-    # We use middleware.set_active_instance to persist the selection.
-    # The session key is an internal detail but useful for debugging response.
     await middleware.set_active_instance(ctx, resolved.id)
-    session_key = await middleware.get_session_key(ctx)
 
     return {
         "success": True,
         "message": f"Active instance set to {resolved.id}",
         "data": {
             "instance": resolved.id,
-            "session_key": session_key,
+            "session_id": getattr(ctx, "session_id", None),
         },
     }

Server/src/transport/legacy/unity_connection.py

@@ -550,17 +550,21 @@ def _resolve_instance_id(self, instance_identifier: str | None, instances: list[
             if self._default_instance_id:
                 instance_identifier = self._default_instance_id
                 logger.debug(f"Using default instance: {instance_identifier}")
+            elif len(instances) == 1:
+                # Sole instance: unambiguous, select it without requiring a hint.
+                return instances[0]
             else:
-                # Use the most recently active instance
-                # Instances with no heartbeat (None) should be sorted last (use 0 as sentinel)
-                sorted_instances = sorted(
-                    instances,
-                    key=lambda inst: inst.last_heartbeat.timestamp() if inst.last_heartbeat else 0.0,
-                    reverse=True,
+                # 2+ instances connected and nothing pinned. Refuse to guess —
+                # silently routing to the most-recently-heartbeated editor lets
+                # an unbound session retarget another project's Unity (#1023).
+                # Mirror the HTTP "multiple connected, no active set" guard.
+                available_ids = [inst.id for inst in instances]
+                raise ConnectionError(
+                    "Multiple Unity instances are connected and none is selected. "
+                    "Pass unity_instance on the call or use set_active_instance "
+                    f"with one of: {available_ids}. "
+                    "Read mcpforunity://instances for current sessions."
                 )
-                logger.info(
-                    f"No instance specified, using most recent: {sorted_instances[0].id}")
-                return sorted_instances[0]
 
         identifier = instance_identifier.strip()
 

Server/src/transport/unity_instance_middleware.py

@@ -55,10 +55,11 @@ class UnityInstanceMiddleware(Middleware):
     for all tool and resource calls.
     """
 
+    # Key used in FastMCP's session-scoped state store for the active instance.
+    _ACTIVE_INSTANCE_STATE_KEY = "mcpforunity.active_instance"
+
     def __init__(self):
         super().__init__()
-        self._active_by_key: dict[str, str] = {}
-        self._lock = RLock()
         self._metadata_lock = RLock()
         self._unity_managed_tool_names: set[str] = set()
         self._tool_alias_to_unity_target: dict[str, str] = {}
@@ -68,43 +69,30 @@ def __init__(self):
         self._tool_visibility_refresh_interval_seconds = 0.5
         self._has_logged_empty_registry_warning = False
 
-    async def get_session_key(self, ctx) -> str:
-        """
-        Derive a stable key for the calling session.
+    async def set_active_instance(self, ctx, instance_id: str) -> None:
+        """Store the active instance for this MCP session.
 
-        Prioritizes client_id for stability.
-        In remote-hosted mode, falls back to user_id for session isolation.
-        Otherwise falls back to 'global' (assuming single-user local mode).
+        Persisted via FastMCP's session-scoped state store, which keys by
+        ``ctx.session_id`` (the MCP-Session-Id header on HTTP, a per-subprocess
+        UUID on stdio). Two MCP sessions cannot share state — see #1023 for the
+        bug this replaces, which keyed on the peer-supplied ``client_id`` and
+        collapsed multiple clients onto the same record.
         """
-        client_id = getattr(ctx, "client_id", None)
-        if isinstance(client_id, str) and client_id:
-            return client_id
-
-        # In remote-hosted mode, use user_id so different users get isolated instance selections
-        user_id = await ctx.get_state("user_id")
-        if isinstance(user_id, str) and user_id:
-            return f"user:{user_id}"
-
-        # Fallback to global for local dev stability
-        return "global"
-
-    async def set_active_instance(self, ctx, instance_id: str) -> None:
-        """Store the active instance for this session."""
-        key = await self.get_session_key(ctx)
-        with self._lock:
-            self._active_by_key[key] = instance_id
+        await ctx.set_state(self._ACTIVE_INSTANCE_STATE_KEY, instance_id)
 
     async def get_active_instance(self, ctx) -> str | None:
-        """Retrieve the active instance for this session."""
-        key = await self.get_session_key(ctx)
-        with self._lock:
-            return self._active_by_key.get(key)
+        """Retrieve the active instance for this MCP session."""
+        return await ctx.get_state(self._ACTIVE_INSTANCE_STATE_KEY)
 
     async def clear_active_instance(self, ctx) -> None:
-        """Clear the stored instance for this session."""
-        key = await self.get_session_key(ctx)
-        with self._lock:
-            self._active_by_key.pop(key, None)
+        """Clear the stored instance for this MCP session.
+
+        Overwrites with None rather than calling ``delete_state``: the read
+        path already treats None as "no active instance", and this keeps the
+        method usable from minimal context shims that don't implement
+        ``delete_state``.
+        """
+        await ctx.set_state(self._ACTIVE_INSTANCE_STATE_KEY, None)
 
     async def _discover_instances(self, ctx) -> list:
         """

Server/tests/integration/test_instance_routing_comprehensive.py

@@ -21,103 +21,86 @@
 from transport.models import SessionList, SessionDetails
 
 
+def _make_stateful_ctx(session_id: str) -> Mock:
+    """Build a Context shim whose set/get/delete_state share one dict.
+
+    The middleware now defers persistence to FastMCP's session-scoped state
+    store, so a useful unit-test mock must actually round-trip values rather
+    than returning fresh Mocks. This helper does that minimally.
+    """
+    state: dict[str, object] = {}
+    ctx = Mock(spec=Context)
+    ctx.session_id = session_id
+    ctx.set_state = AsyncMock(side_effect=lambda k, v: state.__setitem__(k, v))
+    ctx.get_state = AsyncMock(side_effect=lambda k: state.get(k))
+    ctx.delete_state = AsyncMock(side_effect=lambda k: state.pop(k, None))
+    return ctx
+
+
 class TestInstanceRoutingBasics:
     """Test basic middleware functionality."""
 
     @pytest.mark.asyncio
     async def test_middleware_stores_and_retrieves_instance(self):
         """Middleware should store and retrieve instance per session."""
         middleware = UnityInstanceMiddleware()
-        ctx = Mock(spec=Context)
-        ctx.session_id = "test-session-1"
-        ctx.client_id = "test-client-1"
+        ctx = _make_stateful_ctx("test-session-1")
 
-        # Set active instance
         await middleware.set_active_instance(ctx, "TestProject@abc123")
 
-        # Retrieve should return same instance
         assert await middleware.get_active_instance(ctx) == "TestProject@abc123"
 
     @pytest.mark.asyncio
     async def test_middleware_isolates_sessions(self):
-        """Different sessions should have independent instance selections."""
-        middleware = UnityInstanceMiddleware()
+        """Two MCP sessions must not see each other's active instance.
 
-        ctx1 = Mock(spec=Context)
-        ctx1.session_id = "session-1"
-        ctx1.client_id = "client-1"
+        Regression test for #1023: the prior implementation keyed on the
+        peer-supplied client_id and collapsed multiple clients onto a shared
+        record. Each ctx in this test holds its own private state dict, so
+        leakage would surface as a cross-read.
+        """
+        middleware = UnityInstanceMiddleware()
 
-        ctx2 = Mock(spec=Context)
-        ctx2.session_id = "session-2"
-        ctx2.client_id = "client-2"
+        ctx1 = _make_stateful_ctx("session-1")
+        ctx2 = _make_stateful_ctx("session-2")
 
-        # Set different instances for different sessions
         await middleware.set_active_instance(ctx1, "Project1@aaa")
         await middleware.set_active_instance(ctx2, "Project2@bbb")
 
-        # Each session should retrieve its own instance
         assert await middleware.get_active_instance(ctx1) == "Project1@aaa"
         assert await middleware.get_active_instance(ctx2) == "Project2@bbb"
 
-    @pytest.mark.asyncio
-    async def test_middleware_fallback_to_client_id(self):
-        """When session_id unavailable, should use client_id."""
-        middleware = UnityInstanceMiddleware()
-
-        ctx = Mock(spec=Context)
-        ctx.session_id = None
-        ctx.client_id = "client-123"
-
-        await middleware.set_active_instance(ctx, "Project@xyz")
-        assert await middleware.get_active_instance(ctx) == "Project@xyz"
-
-    @pytest.mark.asyncio
-    async def test_middleware_fallback_to_global(self):
-        """When no session/client id, should use 'global' key."""
-        middleware = UnityInstanceMiddleware()
-
-        ctx = Mock(spec=Context)
-        ctx.session_id = None
-        ctx.client_id = None
-        ctx.get_state = AsyncMock(return_value=None)
-
-        await middleware.set_active_instance(ctx, "Project@global")
-        assert await middleware.get_active_instance(ctx) == "Project@global"
-
 
 class TestInstanceRoutingIntegration:
     """Test that instance routing works end-to-end for all tool categories."""
 
     @pytest.mark.asyncio
     async def test_middleware_injects_state_into_context(self):
-        """Middleware on_call_tool should inject instance into ctx state."""
+        """Middleware on_call_tool should inject instance into ctx state.
+
+        After this PR the middleware writes two distinct keys: a persistence
+        key (``mcpforunity.active_instance``) when ``set_active_instance`` is
+        called, and a per-request injection key (``unity_instance``) inside
+        ``on_call_tool``. Tools downstream read ``unity_instance``.
+        """
         middleware = UnityInstanceMiddleware()
 
-        # Create mock context with state management
-        ctx = Mock(spec=Context)
-        ctx.session_id = "test-session"
-        state_storage = {}
-        ctx.set_state = AsyncMock(side_effect=lambda k,
-                             v: state_storage.__setitem__(k, v))
-        ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
+        ctx = _make_stateful_ctx("test-session")
 
-        # Create middleware context
         middleware_ctx = Mock()
         middleware_ctx.fastmcp_context = ctx
 
-        # Set active instance
         await middleware.set_active_instance(ctx, "TestProject@abc123")
 
-        # Mock call_next
         async def mock_call_next(ctx):
             return {"success": True}
 
-        # Execute middleware
         await middleware.on_call_tool(middleware_ctx, mock_call_next)
 
-        # Verify state was injected
-        ctx.set_state.assert_called_once_with(
-            "unity_instance", "TestProject@abc123")
+        # The per-request injection must have happened so tools downstream
+        # can read it; the persistence write is verified by the round-trip
+        # test in TestInstanceRoutingBasics above.
+        ctx.set_state.assert_any_call("unity_instance", "TestProject@abc123")
 
     @pytest.mark.asyncio
     async def test_get_unity_instance_from_context_checks_state(self):

Server/tests/integration/test_middleware_auth_integration.py

@@ -66,38 +66,6 @@ async def test_sets_user_id_in_context_state(self, monkeypatch):
         assert await ctx.get_state("user_id") == "user-55"
 
 
-class TestMiddlewareSessionKey:
-    @pytest.mark.asyncio
-    async def test_get_session_key_uses_user_id_fallback(self):
-        """When no client_id, middleware should use user:$user_id as session key."""
-        from transport.unity_instance_middleware import UnityInstanceMiddleware
-
-        middleware = UnityInstanceMiddleware()
-
-        ctx = DummyContext()
-        # Simulate no client_id attribute
-        if hasattr(ctx, "client_id"):
-            delattr(ctx, "client_id")
-        await ctx.set_state("user_id", "user-77")
-
-        key = await middleware.get_session_key(ctx)
-        assert key == "user:user-77"
-
-    @pytest.mark.asyncio
-    async def test_get_session_key_prefers_client_id(self):
-        """client_id should take precedence over user_id."""
-        from transport.unity_instance_middleware import UnityInstanceMiddleware
-
-        middleware = UnityInstanceMiddleware()
-
-        ctx = DummyContext()
-        ctx.client_id = "client-abc"
-        await ctx.set_state("user_id", "user-77")
-
-        key = await middleware.get_session_key(ctx)
-        assert key == "client-abc"
-
-
 class TestAutoSelectDisabledRemoteHosted:
     @pytest.mark.asyncio
     async def test_auto_select_returns_none_in_remote_hosted(self, monkeypatch):